{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.12.0"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":118765,"databundleVersionId":15231210,"sourceType":"competition"},{"sourceId":14641033,"sourceType":"datasetVersion","datasetId":9349340}],"dockerImageVersionId":31260,"isInternetEnabled":false,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# NullivaRNA v3 V10 - Final Submission\n\n**IMPORTANT:** Please manually attach the competition data:\n1. Click 'Add Data' on the right panel\n2. Search 'stanford-rna-3d-folding-2'\n3. Click 'Add' to attach competition data\n\n**Model:** NullivaRNA v3 (Nullivance-enhanced RNA 3D Predictor)\n**Weights:** Val Loss = 16.02","metadata":{}},{"cell_type":"code","source":"import os\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom tqdm import tqdm\n\n# Check available data sources\nprint('=== Available Data Sources ===')\nif os.path.exists('/kaggle/input'):\n    print(os.listdir('/kaggle/input'))\nelse:\n    print('No /kaggle/input found')\n\n# Find competition data path\nKAGGLE_INPUT = None\nfor name in ['stanford-rna-3d-folding-2', 'stanford-rna-3d-folding', 'rna-3d-folding']:\n    path = f'/kaggle/input/{name}'\n    if os.path.exists(path):\n        KAGGLE_INPUT = path\n        break\n\nif KAGGLE_INPUT is None:\n    # Try to find any directory with test_sequences.csv\n    for d in os.listdir('/kaggle/input'):\n        p = f'/kaggle/input/{d}'\n        if os.path.isdir(p) and os.path.exists(f'{p}/test_sequences.csv'):\n            KAGGLE_INPUT = p\n            break\n\nif KAGGLE_INPUT is None:\n    raise FileNotFoundError('Competition data not found. Please attach stanford-rna-3d-folding-2 dataset manually.')\n\nprint(f'Using: {KAGGLE_INPUT}')\nprint('Files:', os.listdir(KAGGLE_INPUT)[:5])\n\n# Weights path\nWEIGHTS_PATH = '/kaggle/input/nullivarna-full-weights/best_model.pt'\nprint(f'Weights exists: {os.path.exists(WEIGHTS_PATH)}')\n\nMAX_LEN = 384\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f'Device: {device}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === NULLIVANCE MODULES ===\nclass PhiStructure(nn.Module):\n    def __init__(self, eps=1e-8): super().__init__(); self.eps = eps\n    def entropy(self, p, dim=-1): return -(p.clamp(min=self.eps) * torch.log(p.clamp(min=self.eps))).sum(dim=dim)\n    def forward(self, attn): return (1.0 - self.entropy(attn, dim=-1) / np.log(max(attn.shape[-1], 2))).mean()\n    def loss(self, attn): return (1.0 - self.forward(attn)).clamp(0, 1)\n\nclass GammaCompatibility(nn.Module):\n    def __init__(self): super().__init__(); self.scale = nn.Parameter(torch.tensor(0.3))\n    def forward(self, seq, msa=None):\n        B, L = seq.shape; dev = seq.device\n        pair = torch.zeros(6, 6, device=dev)\n        pair[0,1]=pair[1,0]=1.0; pair[2,3]=pair[3,2]=1.0; pair[2,1]=pair[1,2]=0.5\n        s = seq.clamp(0, 5)\n        bias = pair[s.unsqueeze(-1).expand(B,L,L), s.unsqueeze(-2).expand(B,L,L)]\n        mask = (torch.abs(torch.arange(L,device=dev).unsqueeze(0)-torch.arange(L,device=dev).unsqueeze(1))>=4)\n        return self.scale * bias * mask.unsqueeze(0).float()\n\nclass PsiTension(nn.Module):\n    def __init__(self, d): super().__init__(); self.proj = nn.Sequential(nn.LayerNorm(d), nn.Linear(d,1), nn.Sigmoid())\n    def forward(self, h, attn=None): return self.proj(h).squeeze(-1) * 0.1\n\nclass OPlusAggregation(nn.Module):\n    def forward(self, pred, w=None): return pred.mean(dim=1)\n\nclass WaveFrameEmbedding(nn.Module):\n    def __init__(self, d, max_len=2048):\n        super().__init__(); self.freq = nn.Parameter(torch.exp(torch.linspace(0,-4,d//2)))\n    def forward(self, L, dev):\n        pos = torch.arange(L,device=dev).float()/2048\n        omega = self.freq.to(dev) * pos.unsqueeze(-1) * 2*np.pi*10\n        return torch.cat([torch.sin(omega), torch.cos(omega)], dim=-1).unsqueeze(0)\n\nprint('Nullivance modules OK')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === MODEL ===\nclass NullivaRNAConfig:\n    def __init__(self, d_model=256, n_layers=6, n_heads=8, d_ff=1024, max_len=384, dropout=0.1, n_preds=5):\n        self.d_model=d_model; self.n_layers=n_layers; self.n_heads=n_heads; self.d_ff=d_ff\n        self.max_len=max_len; self.dropout=dropout; self.n_preds=n_preds; self.d_head=d_model//n_heads\n\nclass RNAEmbedding(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.emb = nn.Embedding(6, cfg.d_model, padding_idx=4)\n        self.pos = WaveFrameEmbedding(cfg.d_model, cfg.max_len)\n        self.norm = nn.LayerNorm(cfg.d_model); self.drop = nn.Dropout(cfg.dropout)\n    def forward(self, seq): return self.drop(self.norm(self.emb(seq) + self.pos(seq.shape[1], seq.device)))\n\nclass GammaAttn(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.H=cfg.n_heads; self.dh=cfg.d_head\n        self.qkv=nn.Linear(cfg.d_model, 3*cfg.d_model); self.out=nn.Linear(cfg.d_model, cfg.d_model)\n        self.g_scale=nn.Parameter(torch.tensor(0.1)); self.drop=nn.Dropout(cfg.dropout)\n    def forward(self, x, g_bias=None, ret_attn=False):\n        B,L,D=x.shape\n        qkv=self.qkv(x).view(B,L,3,self.H,self.dh).permute(2,0,3,1,4)\n        Q,K,V=qkv[0],qkv[1],qkv[2]\n        sc = Q@K.transpose(-2,-1)/np.sqrt(self.dh)\n        if g_bias is not None: sc = sc + self.g_scale*g_bias.unsqueeze(1)\n        attn = self.drop(F.softmax(sc, dim=-1))\n        out = (attn@V).transpose(1,2).reshape(B,L,D)\n        return (self.out(out), attn) if ret_attn else (self.out(out), None)\n\nclass TransBlock(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.attn=GammaAttn(cfg); self.n1=nn.LayerNorm(cfg.d_model)\n        self.ff=nn.Sequential(nn.Linear(cfg.d_model,cfg.d_ff),nn.GELU(),nn.Dropout(cfg.dropout),nn.Linear(cfg.d_ff,cfg.d_model),nn.Dropout(cfg.dropout))\n        self.n2=nn.LayerNorm(cfg.d_model)\n    def forward(self, x, g=None, ret=False):\n        a,w = self.attn(self.n1(x), g, ret)\n        return x+a+self.ff(self.n2(x+a)), w\n\nclass IPA(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.q=nn.Linear(cfg.d_model,cfg.d_model); self.k=nn.Linear(cfg.d_model,cfg.d_model)\n        self.v=nn.Linear(cfg.d_model,cfg.d_model); self.out=nn.Linear(cfg.d_model,cfg.d_model)\n        self.coord=nn.Linear(cfg.d_model,3); self.norm=nn.LayerNorm(cfg.d_model)\n    def forward(self, f, c):\n        B,L,D=f.shape\n        sc = self.q(f)@self.k(f).transpose(-2,-1)/np.sqrt(D)\n        dist = torch.norm(c.unsqueeze(2)-c.unsqueeze(1), dim=-1).clamp(min=0.1)\n        agg = F.softmax(sc - 0.05*dist, dim=-1)@self.v(f)\n        up_f = f + self.out(self.norm(agg))\n        return up_f, c + self.coord(up_f)\n\nclass DivHeads(nn.Module):\n    def __init__(self, cfg):\n        super().__init__()\n        self.heads = nn.ModuleList([nn.Linear(cfg.d_model,3) for _ in range(cfg.n_preds)])\n        self.psi = PsiTension(cfg.d_model)\n    def forward(self, f, base, attn=None):\n        psi = self.psi(f, attn)\n        return torch.stack([base + h(f) + (torch.randn_like(base)*psi.unsqueeze(-1)*0.05 if i > 0 else 0) for i,h in enumerate(self.heads)], dim=1)\n\nclass NullivaRNAv3(nn.Module):\n    def __init__(self, cfg=None):\n        super().__init__()\n        self.cfg = cfg or NullivaRNAConfig()\n        self.emb = RNAEmbedding(self.cfg)\n        self.gamma = GammaCompatibility()\n        self.layers = nn.ModuleList([TransBlock(self.cfg) for _ in range(self.cfg.n_layers)])\n        self.ipa = IPA(self.cfg)\n        self.init_coords = nn.Linear(self.cfg.d_model, 3)\n        self.div = DivHeads(self.cfg)\n        self.oplus = OPlusAggregation()\n        self.phi = PhiStructure()\n    def forward(self, seq, ret_all=False):\n        B,L = seq.shape; dev = seq.device\n        x = self.emb(seq); g = self.gamma(seq)\n        attns = []\n        for layer in self.layers:\n            x, a = layer(x, g, ret=True)\n            if a is not None: attns.append(a)\n        init_c = self.init_coords(x)\n        x, ref_c = self.ipa(x, init_c)\n        div_c = self.div(x, ref_c, attns[-1] if attns else None)\n        return {'all_coords': div_c}\n\nprint('Model OK')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === LOAD WEIGHTS ===\ncfg = NullivaRNAConfig(d_model=256, n_layers=6, n_heads=8, d_ff=1024, max_len=MAX_LEN)\nmodel = NullivaRNAv3(cfg).to(device)\n\nif os.path.exists(WEIGHTS_PATH):\n    print('Loading trained weights...')\n    state_dict = torch.load(WEIGHTS_PATH, map_location=device)\n    model.load_state_dict(state_dict)\nelse:\n    print('WARNING: No weights found, using random init')\n\nmodel.eval()\nprint(f'Model ready: {sum(p.numel() for p in model.parameters()):,} params')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === INFERENCE ===\ntest_df = pd.read_csv(f'{KAGGLE_INPUT}/test_sequences.csv')\nprint(f'Test: {len(test_df)} sequences')\n\nMAP = {'A':0,'U':1,'G':2,'C':3,'T':1,'N':5}\nrows = []\n\nwith torch.no_grad():\n    for _, row in tqdm(test_df.iterrows(), total=len(test_df)):\n        tid = str(row['target_id'])\n        seq = str(row['sequence']).upper().replace('T','U')\n        enc = [MAP.get(b,5) for b in seq[:MAX_LEN]]\n        if len(enc) < MAX_LEN: enc += [4]*(MAX_LEN-len(enc))\n        t = torch.tensor(enc, dtype=torch.long).unsqueeze(0).to(device)\n        coords = model(t)['all_coords'].cpu().numpy()[0]\n        for i in range(min(len(seq), MAX_LEN)):\n            rows.append([f'{tid}_{i+1}', seq[i], i+1] + coords[:, i, :].flatten().tolist())\n\nprint(f'Rows: {len(rows)}')","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# === SAVE submission.csv ===\ncols = ['ID', 'resname', 'resid']\nfor k in range(1, 6): cols += [f'x_{k}', f'y_{k}', f'z_{k}']\n\ndf = pd.DataFrame(rows, columns=cols)\ndf.to_csv('/kaggle/working/submission.csv', index=False)\n\nprint(f'Shape: {df.shape}')\nprint(f'Columns: {list(df.columns)}')\nprint(df.head())\nprint()\nprint('SAVED: /kaggle/working/submission.csv')","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}