{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.11.13","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python Docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load\n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the read-only \"../input/\" directory\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        print(os.path.join(dirname, filename))\n\n# You can write up to 20GB to the current directory (/kaggle/working/) that gets preserved as output when you create a version using \"Save & Run All\" \n# You can also write temporary files to /kaggle/temp/, but they won't be saved outside of the current session","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install obonet Bio\n!pip install biopython pandas scipy numpy transformers \n!pip install torch torchvision torchaudio","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T05:36:05.989590Z","iopub.execute_input":"2025-12-23T05:36:05.990425Z","iopub.status.idle":"2025-12-23T05:37:37.072912Z","shell.execute_reply.started":"2025-12-23T05:36:05.990394Z","shell.execute_reply":"2025-12-23T05:37:37.072147Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!nvidia-smi","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T05:37:37.074772Z","iopub.execute_input":"2025-12-23T05:37:37.075066Z","iopub.status.idle":"2025-12-23T05:37:37.307426Z","shell.execute_reply.started":"2025-12-23T05:37:37.075038Z","shell.execute_reply":"2025-12-23T05:37:37.306636Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nprint(torch.__version__)\nprint(torch.cuda.is_available())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T05:37:37.308388Z","iopub.execute_input":"2025-12-23T05:37:37.308651Z","iopub.status.idle":"2025-12-23T05:37:37.397858Z","shell.execute_reply.started":"2025-12-23T05:37:37.308625Z","shell.execute_reply":"2025-12-23T05:37:37.397172Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# FIXED & TUNED PIPELINE (Run This)\n# ==========================================\nimport os\n# FIX 1: Prevent memory fragmentation\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nfrom transformers import AutoTokenizer, EsmModel\nfrom torch.utils.data import TensorDataset, DataLoader\nimport obonet\nimport networkx as nx\nfrom collections import defaultdict, Counter\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score\n\n# -------------------------\n# CONFIGURATION\n# -------------------------\nclass CFG:\n    # PATHS\n    BASE = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    if os.path.exists(f\"{BASE}/Train\"):\n        TRAIN_FASTA = f\"{BASE}/Train/train_sequences.fasta\"\n        TRAIN_TERMS = f\"{BASE}/Train/train_terms.tsv\"\n        OBO = f\"{BASE}/Train/go-basic.obo\"\n    else:\n        TRAIN_FASTA = f\"{BASE}/train_sequences.fasta\"\n        TRAIN_TERMS = f\"{BASE}/train_terms.tsv\"\n        OBO = f\"{BASE}/go-basic.obo\"\n        \n    TEST_FASTA  = f\"{BASE}/Test/testsuperset.fasta\" if os.path.exists(f\"{BASE}/Test\") else f\"{BASE}/testsuperset.fasta\"\n    OUT_DIR = \"/kaggle/working\" \n\n    # MODEL\n    MODEL = \"facebook/esm2_t12_35M_UR50D\" \n    EMBED = 480   \n    MAX_LEN = 1022        \n\n    # --- FIX 2: LOWER BATCH SIZE ---\n    # 256 crashed -> 64 is safe (32 sequences per GPU)\n    BS_EMBED = 64        \n    \n    BS_TRAIN = 2048 # Training is cheap, keep this high\n    LR = 1e-3\n    EPOCHS = 12\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    MAX_LABELS = { \"MFO\": 2000, \"BPO\": 4000, \"CCO\": 2000 }\n\n# -------------------------\n# 1. ONTOLOGY (CPU)\n# -------------------------\nclass Ontology:\n    def __init__(self, obo):\n        print(\">> Loading OBO Ontology...\")\n        self.graph = obonet.read_obo(obo)\n        self.ancestors = {}\n        for node in tqdm(self.graph.nodes, desc=\"Indexing Ancestors\"):\n            self.ancestors[node] = nx.ancestors(self.graph, node)\n        \n        self.ns_map = {\n            \"molecular_function\": \"MFO\",\n            \"biological_process\": \"BPO\",\n            \"cellular_component\": \"CCO\"\n        }\n\n    def get_namespace(self, term):\n        if term in self.graph.nodes:\n            ns = self.graph.nodes[term].get(\"namespace\")\n            return self.ns_map.get(ns, None)\n        return None\n\ndef load_labels_optimized(path, ont, aspect, max_labels):\n    print(f\"Loading labels for {aspect}...\")\n    df = pd.read_csv(path, sep=\"\\t\", names=[\"pid\", \"term\", \"aspect\"], header=0, usecols=[0, 1, 2])\n    # Pre-filter by string match first (faster)\n    if aspect == \"MFO\": df = df[df['aspect'] == 'molecular_function']\n    elif aspect == \"BPO\": df = df[df['aspect'] == 'biological_process']\n    elif aspect == \"CCO\": df = df[df['aspect'] == 'cellular_component']\n    \n    term_counts = Counter()\n    pid_to_terms = defaultdict(set)\n    \n    print(f\"   Propagating Terms...\")\n    for pid, term in tqdm(zip(df['pid'], df['term']), total=len(df)):\n        to_add = {term}\n        if term in ont.ancestors:\n            to_add.update([a for a in ont.ancestors[term] if ont.get_namespace(a) == aspect])\n        pid_to_terms[pid].update(to_add)\n        term_counts.update(to_add)\n        \n    top_terms = [t for t, _ in term_counts.most_common(max_labels)]\n    term_to_idx = {t: i for i, t in enumerate(top_terms)}\n    \n    final_labels = {}\n    valid_pids = []\n    \n    for pid, terms in pid_to_terms.items():\n        indices = [term_to_idx[t] for t in terms if t in term_to_idx]\n        if indices:\n            vec = np.zeros(max_labels, dtype=np.float32)\n            vec[indices] = 1.0\n            final_labels[pid] = vec\n            valid_pids.append(pid)\n            \n    del df, term_counts, pid_to_terms\n    gc.collect()\n    return valid_pids, final_labels, top_terms\n\n# -------------------------\n# 2. EMBEDDING (GPU)\n# -------------------------\ndef clean_id(header):\n    if \"|\" in header: return header.split(\"|\")[1]\n    return header\n\ndef generate_embeddings(fasta_path, name_prefix):\n    npy_path = os.path.join(CFG.OUT_DIR, f\"{name_prefix}_embs.npy\")\n    ids_path = os.path.join(CFG.OUT_DIR, f\"{name_prefix}_ids.npy\")\n    \n    if os.path.exists(npy_path) and os.path.exists(ids_path):\n        print(f\">> Found cached {name_prefix}. Skipping.\")\n        return ids_path, npy_path\n\n    print(f\">> Extracting {name_prefix}...\")\n    \n    # Init Model\n    tokenizer = AutoTokenizer.from_pretrained(CFG.MODEL)\n    base_model = EsmModel.from_pretrained(CFG.MODEL)\n    \n    if torch.cuda.device_count() > 1:\n        print(f\"   Using {torch.cuda.device_count()} GPUs\")\n        model = nn.DataParallel(base_model)\n    else:\n        model = base_model\n    \n    model = model.to(CFG.DEVICE)\n    model.eval()\n\n    # Read\n    ids, seqs = [], []\n    for rec in SeqIO.parse(fasta_path, \"fasta\"):\n        ids.append(clean_id(rec.id))\n        seqs.append(str(rec.seq)[:CFG.MAX_LEN])\n    \n    # Sort\n    perm_idx = np.argsort([len(s) for s in seqs])\n    ids = [ids[i] for i in perm_idx]\n    seqs = [seqs[i] for i in perm_idx]\n    \n    embeddings = []\n    \n    # Inference\n    with torch.no_grad(), torch.amp.autocast('cuda'):\n        for i in tqdm(range(0, len(seqs), CFG.BS_EMBED)):\n            batch = seqs[i:i + CFG.BS_EMBED]\n            toks = tokenizer(batch, return_tensors=\"pt\", padding=True, truncation=True, max_length=CFG.MAX_LEN).to(CFG.DEVICE)\n            out = model(**toks).last_hidden_state[:, 0, :]\n            embeddings.append(out.cpu().numpy().astype(np.float32))\n            \n    full_embs = np.vstack(embeddings)\n    \n    np.save(npy_path, full_embs)\n    np.save(ids_path, ids)\n    \n    del model, base_model, tokenizer, embeddings, full_embs, seqs\n    torch.cuda.empty_cache()\n    gc.collect()\n    \n    return ids_path, npy_path\n\n# -------------------------\n# 3. TRAINING UTILS\n# -------------------------\nclass MLP(nn.Module):\n    def __init__(self, in_dim, out_dim):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(in_dim, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.5),\n            nn.Linear(512, out_dim)\n        )\n    def forward(self, x): return self.net(x)\n\ndef find_best_threshold(y_true, y_pred_probs):\n    best_f1 = 0\n    best_th = 0.2\n    for th in np.arange(0.1, 0.6, 0.05):\n        y_pred = (y_pred_probs > th).astype(int)\n        score = f1_score(y_true, y_pred, average='micro')\n        if score > best_f1:\n            best_f1 = score\n            best_th = th\n    return best_th, best_f1\n\n# -------------------------\n# 4. MAIN\n# -------------------------\ndef main():\n    # 1. Embed (Disk IO)\n    tr_ids_path, tr_npy_path = generate_embeddings(CFG.TRAIN_FASTA, \"train\")\n    te_ids_path, te_npy_path = generate_embeddings(CFG.TEST_FASTA, \"test\")\n    \n    # 2. Load IDs\n    tr_ids = np.load(tr_ids_path)\n    te_ids = np.load(te_ids_path)\n    tr_id_map = {pid: i for i, pid in enumerate(tr_ids)}\n    \n    # 3. Memory Map Embeddings\n    Xtr_disk = np.load(tr_npy_path, mmap_mode='r')\n    Xte_disk = np.load(te_npy_path, mmap_mode='r')\n    \n    ont = Ontology(CFG.OBO)\n    submission_lines = []\n    \n    for aspect in [\"MFO\", \"BPO\", \"CCO\"]:\n        print(f\"\\n=== Processing {aspect} ===\")\n        valid_pids, labels_dict, term_list = load_labels_optimized(CFG.TRAIN_TERMS, ont, aspect, CFG.MAX_LABELS[aspect])\n        \n        # Gather Indices\n        indices = []\n        Y_local = []\n        for pid in valid_pids:\n            if pid in tr_id_map:\n                indices.append(tr_id_map[pid])\n                Y_local.append(labels_dict[pid])\n        \n        if not indices: continue\n        \n        # Load Subset to RAM\n        print(f\"   Loading {len(indices)} samples to RAM...\")\n        X_local = Xtr_disk[indices]\n        Y_local = np.array(Y_local)\n        \n        X_train, X_val, y_train, y_val = train_test_split(X_local, Y_local, test_size=0.1, random_state=42)\n        del X_local, Y_local\n        gc.collect()\n        \n        # Train\n        train_ds = TensorDataset(torch.tensor(X_train), torch.tensor(y_train))\n        train_loader = DataLoader(train_ds, batch_size=CFG.BS_TRAIN, shuffle=True)\n        \n        model = MLP(CFG.EMBED, len(term_list)).to(CFG.DEVICE)\n        optimizer = optim.Adam(model.parameters(), lr=CFG.LR)\n        criterion = nn.BCEWithLogitsLoss()\n        \n        print(f\"   Training...\")\n        for epoch in range(CFG.EPOCHS):\n            model.train()\n            for x, y in train_loader:\n                x, y = x.to(CFG.DEVICE), y.to(CFG.DEVICE)\n                optimizer.zero_grad()\n                out = model(x)\n                loss = criterion(out, y)\n                loss.backward()\n                optimizer.step()\n        \n        # Threshold\n        model.eval()\n        with torch.no_grad():\n            v_x = torch.tensor(X_val).to(CFG.DEVICE)\n            val_probs = torch.sigmoid(model(v_x)).cpu().numpy()\n        best_th, val_f1 = find_best_threshold(y_val, val_probs)\n        print(f\"   Best Threshold: {best_th:.2f} (F1: {val_f1:.4f})\")\n        \n        # Predict on Test (Chunked)\n        print(\"   Predicting...\")\n        te_preds = []\n        total_test = len(Xte_disk)\n        with torch.no_grad():\n            for i in range(0, total_test, CFG.BS_TRAIN):\n                batch = Xte_disk[i:i+CFG.BS_TRAIN]\n                batch_t = torch.tensor(batch).to(CFG.DEVICE)\n                out = torch.sigmoid(model(batch_t)).cpu().numpy()\n                te_preds.append(out)\n        te_preds = np.vstack(te_preds)\n        \n        # Format\n        for i, pid in enumerate(te_ids):\n            probs = te_preds[i]\n            idxs = np.where(probs >= best_th)[0]\n            if len(idxs) == 0: idxs = np.argsort(probs)[-2:]\n            for idx in idxs:\n                submission_lines.append(f\"{pid}\\t{term_list[idx]}\\t{probs[idx]:.3f}\")\n                \n        del model, X_train, X_val, train_loader\n        torch.cuda.empty_cache()\n        gc.collect()\n\n    print(\">> Writing Submission...\")\n    with open(\"submission.tsv\", \"w\") as f:\n        f.write(\"\\n\".join(submission_lines))\n    print(\">> DONE.\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# CAFA-6 FIXED & WORKING PIPELINE\n# ==========================================\nimport os\nos.environ[\"PYTORCH_CUDA_ALLOC_CONF\"] = \"expandable_segments:True\"\n\nimport gc\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nimport numpy as np\nimport pandas as pd\nfrom tqdm import tqdm\nfrom Bio import SeqIO\nfrom transformers import AutoTokenizer, EsmModel\nfrom torch.utils.data import TensorDataset, DataLoader\nimport obonet\nimport networkx as nx\nfrom collections import defaultdict, Counter\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score\n\n# -------------------------\n# CONFIG\n# -------------------------\nclass CFG:\n    BASE = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    TRAIN_FASTA = f\"{BASE}/Train/train_sequences.fasta\"\n    TRAIN_TERMS = f\"{BASE}/Train/train_terms.tsv\"\n    TEST_FASTA  = f\"{BASE}/Test/testsuperset.fasta\"\n    OBO = f\"{BASE}/Train/go-basic.obo\"\n    OUT_DIR = \"/kaggle/working\"\n\n    MODEL = \"facebook/esm2_t12_35M_UR50D\"\n    EMBED = 480\n    MAX_LEN = 1022\n\n    BS_EMBED = 64\n    BS_TRAIN = 2048\n    LR = 1e-3\n    EPOCHS = 12\n\n    DEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    MAX_LABELS = {\n        \"MFO\": 2000,\n        \"BPO\": 4000,\n        \"CCO\": 2000\n    }\n\n# -------------------------\n# 1. ONTOLOGY\n# -------------------------\nclass Ontology:\n    def __init__(self, obo_path):\n        print(\">> Loading GO ontology\")\n        self.graph = obonet.read_obo(obo_path)\n        self.ancestors = {\n            n: nx.ancestors(self.graph, n)\n            for n in self.graph.nodes\n        }\n        self.ns_map = {\n            \"molecular_function\": \"MFO\",\n            \"biological_process\": \"BPO\",\n            \"cellular_component\": \"CCO\"\n        }\n\n    def namespace(self, term):\n        ns = self.graph.nodes[term].get(\"namespace\")\n        return self.ns_map.get(ns, None)\n\n# -------------------------\n# 2. LOAD + PROPAGATE LABELS (FIXED)\n# -------------------------\ndef load_labels(path, ont, aspect, max_labels):\n    print(f\">> Loading labels for {aspect}\")\n\n    df = pd.read_csv(path, sep=\"\\t\")\n\n    # ✅ CAFA-6 column names\n    df = df.rename(columns={\n        df.columns[0]: \"pid\",\n        \"go_id\": \"term\"\n    })\n\n    # ✅ CAFA-6 namespace mapping\n    aspect_map = {\n        \"MFO\": \"molecular_function\",\n        \"BPO\": \"biological_process\",\n        \"CCO\": \"cellular_component\"\n    }\n\n    df[\"namespace\"] = df[\"namespace\"].astype(str).str.strip().str.lower()\n    target_ns = aspect_map[aspect]\n\n    df = df[df[\"namespace\"] == target_ns]\n\n    print(f\"   Rows after filter: {len(df)}\")\n    if len(df) == 0:\n        print(\"   ❌ No rows — check namespace column\")\n        print(\"   Seen namespaces:\", df[\"namespace\"].unique())\n        return [], {}, []\n\n    pid2terms = defaultdict(set)\n    term_freq = Counter()\n\n    for pid, term in tqdm(zip(df.pid, df.term), total=len(df)):\n        all_terms = {term}\n\n        # ✅ ontology propagation\n        if term in ont.ancestors:\n            for a in ont.ancestors[term]:\n                if ont.namespace(a) == aspect:\n                    all_terms.add(a)\n\n        pid2terms[pid].update(all_terms)\n        term_freq.update(all_terms)\n\n    # ✅ limit label space\n    top_terms = [t for t, _ in term_freq.most_common(max_labels)]\n    t2i = {t: i for i, t in enumerate(top_terms)}\n\n    labels = {}\n    valid_pids = []\n\n    for pid, terms in pid2terms.items():\n        idxs = [t2i[t] for t in terms if t in t2i]\n        if idxs:\n            vec = np.zeros(len(top_terms), dtype=np.float32)\n            vec[idxs] = 1.0\n            labels[pid] = vec\n            valid_pids.append(pid)\n\n    return valid_pids, labels, top_terms\n\n\n# -------------------------\n# 3. EMBEDDINGS\n# -------------------------\ndef clean_id(h):\n    return h.split(\"|\")[1] if \"|\" in h else h\n\ndef generate_embeddings(fasta, name):\n    npy = f\"{CFG.OUT_DIR}/{name}_emb.npy\"\n    idsf = f\"{CFG.OUT_DIR}/{name}_ids.npy\"\n\n    if os.path.exists(npy):\n        return idsf, npy\n\n    tokenizer = AutoTokenizer.from_pretrained(CFG.MODEL)\n    base = EsmModel.from_pretrained(CFG.MODEL)\n    model = nn.DataParallel(base).to(CFG.DEVICE)\n    model.eval()\n\n    ids, seqs = [], []\n    for r in SeqIO.parse(fasta, \"fasta\"):\n        ids.append(clean_id(r.id))\n        seqs.append(str(r.seq)[:CFG.MAX_LEN])\n\n    order = np.argsort([len(s) for s in seqs])\n    ids = [ids[i] for i in order]\n    seqs = [seqs[i] for i in order]\n\n    embs = []\n    with torch.no_grad(), torch.amp.autocast(\"cuda\"):\n        for i in tqdm(range(0, len(seqs), CFG.BS_EMBED)):\n            batch = seqs[i:i+CFG.BS_EMBED]\n            toks = tokenizer(batch, return_tensors=\"pt\",\n                             padding=True, truncation=True).to(CFG.DEVICE)\n            out = model(**toks).last_hidden_state[:, 0, :]\n            embs.append(out.cpu().numpy())\n\n    embs = np.vstack(embs).astype(np.float32)\n    np.save(npy, embs)\n    np.save(idsf, ids)\n    return idsf, npy\n\n# -------------------------\n# 4. CLASSIFIER\n# -------------------------\nclass MLP(nn.Module):\n    def __init__(self, d_in, d_out):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(d_in, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.5),\n            nn.Linear(512, d_out)\n        )\n    def forward(self, x): return self.net(x)\n\ndef best_threshold(y, p):\n    best, th = 0, 0.3\n    for t in np.arange(0.1, 0.6, 0.05):\n        f = f1_score(y, (p > t), average=\"micro\")\n        if f > best:\n            best, th = f, t\n    return th\n\n# -------------------------\n# 5. MAIN\n# -------------------------\ndef main():\n    tr_ids_p, tr_emb_p = generate_embeddings(CFG.TRAIN_FASTA, \"train\")\n    te_ids_p, te_emb_p = generate_embeddings(CFG.TEST_FASTA, \"test\")\n\n    tr_ids = np.load(tr_ids_p)\n    te_ids = np.load(te_ids_p)\n    tr_map = {p: i for i, p in enumerate(tr_ids)}\n\n    Xtr = np.load(tr_emb_p, mmap_mode=\"r\")\n    Xte = np.load(te_emb_p, mmap_mode=\"r\")\n\n    ont = Ontology(CFG.OBO)\n    submission = []\n\n    for aspect in [\"MFO\", \"BPO\", \"CCO\"]:\n        pids, labels, terms = load_labels(CFG.TRAIN_TERMS, ont, aspect, CFG.MAX_LABELS[aspect])\n\n        idx, Y = [], []\n        for p in pids:\n            if p in tr_map:\n                idx.append(tr_map[p])\n                Y.append(labels[p])\n\n        if not idx:\n            continue\n\n        X = Xtr[idx]\n        Y = np.array(Y)\n\n        Xtrn, Xv, Ytrn, Yv = train_test_split(X, Y, test_size=0.1, random_state=42)\n\n        ds = TensorDataset(torch.tensor(Xtrn), torch.tensor(Ytrn))\n        dl = DataLoader(ds, batch_size=CFG.BS_TRAIN, shuffle=True)\n\n        model = MLP(CFG.EMBED, len(terms)).to(CFG.DEVICE)\n        opt = optim.Adam(model.parameters(), lr=CFG.LR)\n        loss_fn = nn.BCEWithLogitsLoss()\n\n        for _ in range(CFG.EPOCHS):\n            for x, y in dl:\n                x, y = x.to(CFG.DEVICE), y.to(CFG.DEVICE)\n                opt.zero_grad()\n                loss = loss_fn(model(x), y)\n                loss.backward()\n                opt.step()\n\n        with torch.no_grad():\n            pv = torch.sigmoid(model(torch.tensor(Xv).to(CFG.DEVICE))).cpu().numpy()\n        th = best_threshold(Yv, pv)\n\n        with torch.no_grad():\n            preds = []\n            for i in range(0, len(Xte), CFG.BS_TRAIN):\n                b = torch.tensor(Xte[i:i+CFG.BS_TRAIN]).to(CFG.DEVICE)\n                preds.append(torch.sigmoid(model(b)).cpu().numpy())\n        preds = np.vstack(preds)\n\n        for i, pid in enumerate(te_ids):\n            idxs = np.where(preds[i] >= th)[0]\n            if len(idxs) == 0:\n                idxs = np.argsort(preds[i])[-2:]\n            for j in idxs:\n                submission.append(f\"{pid}\\t{terms[j]}\\t{preds[i][j]:.3f}\")\n\n    with open(\"submission.tsv\", \"w\") as f:\n        f.write(\"\\n\".join(submission))\n\n    print(\">> SUBMISSION READY\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T05:38:04.494755Z","iopub.execute_input":"2025-12-23T05:38:04.495279Z","iopub.status.idle":"2025-12-23T06:26:12.331397Z","shell.execute_reply.started":"2025-12-23T05:38:04.495254Z","shell.execute_reply":"2025-12-23T06:26:12.330377Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ==========================================\n# CAFA-6 SUBMISSION-ONLY PIPELINE\n# ==========================================\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom tqdm import tqdm\nfrom collections import defaultdict, Counter\nfrom sklearn.model_selection import train_test_split\nfrom sklearn.metrics import f1_score\nimport obonet\nimport networkx as nx\n\n# -------------------------\n# CONFIG\n# -------------------------\nBASE = \"/kaggle/input/cafa-6-protein-function-prediction\"\nWORK = \"/kaggle/working\"\n\nTRAIN_TERMS = f\"{BASE}/Train/train_terms.tsv\"\nOBO = f\"{BASE}/Train/go-basic.obo\"\n\nDEVICE = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\nMAX_LABELS = {\n    \"MFO\": 2000,\n    \"BPO\": 4000,\n    \"CCO\": 2000\n}\n\nBS = 2048\nLR = 1e-3\nEPOCHS = 12\nEMBED = 480\n\n# -------------------------\n# ONTOLOGY\n# -------------------------\nclass Ontology:\n    def __init__(self, obo_path):\n        print(\">> Loading GO ontology\")\n        self.graph = obonet.read_obo(obo_path)\n        self.ancestors = {\n            n: nx.ancestors(self.graph, n)\n            for n in self.graph.nodes\n        }\n        self.ns_map = {\n            \"molecular_function\": \"MFO\",\n            \"biological_process\": \"BPO\",\n            \"cellular_component\": \"CCO\"\n        }\n\n    def namespace(self, term):\n        ns = self.graph.nodes[term].get(\"namespace\")\n        return self.ns_map.get(ns)\n\n# -------------------------\n# LOAD LABELS (FIXED)\n# -------------------------\ndef load_labels(path, ont, aspect, max_labels):\n    print(f\">> Loading labels for {aspect}\")\n\n    df = pd.read_csv(path, sep=\"\\t\")\n\n    # CAFA-6 canonical columns\n    df = df.rename(columns={\n        \"EntryID\": \"pid\",\n        \"term\": \"term\",\n        \"aspect\": \"aspect\"\n    })\n\n    # CAFA-6 aspect mapping (CRITICAL FIX)\n    aspect_code = {\n        \"MFO\": \"F\",\n        \"BPO\": \"P\",\n        \"CCO\": \"C\"\n    }[aspect]\n\n    df[\"aspect\"] = df[\"aspect\"].astype(str).str.strip()\n    df = df[df[\"aspect\"] == aspect_code]\n\n    print(f\"   Rows after filter: {len(df)}\")\n    if len(df) == 0:\n        print(\"   ❌ No rows — aspect mismatch\")\n        return [], {}, []\n\n    pid2terms = defaultdict(set)\n    term_freq = Counter()\n\n    for pid, term in tqdm(zip(df.pid, df.term), total=len(df)):\n        all_terms = {term}\n\n        # ✅ GO ancestor propagation\n        if term in ont.ancestors:\n            for a in ont.ancestors[term]:\n                if ont.namespace(a) == aspect:\n                    all_terms.add(a)\n\n        pid2terms[pid].update(all_terms)\n        term_freq.update(all_terms)\n\n    # Limit label space\n    top_terms = [t for t, _ in term_freq.most_common(max_labels)]\n    t2i = {t: i for i, t in enumerate(top_terms)}\n\n    labels = {}\n    valid_pids = []\n\n    for pid, terms in pid2terms.items():\n        idxs = [t2i[t] for t in terms if t in t2i]\n        if idxs:\n            v = np.zeros(len(top_terms), dtype=np.float32)\n            v[idxs] = 1.0\n            labels[pid] = v\n            valid_pids.append(pid)\n\n    return valid_pids, labels, top_terms\n\n\n# -------------------------\n# MODEL\n# -------------------------\nclass MLP(nn.Module):\n    def __init__(self, d_in, d_out):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Linear(d_in, 512),\n            nn.BatchNorm1d(512),\n            nn.ReLU(),\n            nn.Dropout(0.5),\n            nn.Linear(512, d_out)\n        )\n    def forward(self, x):\n        return self.net(x)\n\ndef best_threshold(y, p):\n    best, th = 0, 0.3\n    for t in np.arange(0.1, 0.6, 0.05):\n        f = f1_score(y, (p > t), average=\"micro\")\n        if f > best:\n            best, th = f, t\n    return th\n\n# -------------------------\n# MAIN\n# -------------------------\ndef main():\n    # ✅ load cached embeddings\n    Xtr = np.load(f\"{WORK}/train_emb.npy\", mmap_mode=\"r\")\n    Xte = np.load(f\"{WORK}/test_emb.npy\")\n    tr_ids = np.load(f\"{WORK}/train_ids.npy\")\n    te_ids = np.load(f\"{WORK}/test_ids.npy\")\n\n    tr_map = {p: i for i, p in enumerate(tr_ids)}\n\n    ont = Ontology(OBO)\n    submission = []\n\n    for aspect in [\"MFO\", \"BPO\", \"CCO\"]:\n        pids, labels, terms = load_labels(TRAIN_TERMS, ont, aspect, MAX_LABELS[aspect])\n\n        idx, Y = [], []\n        for p in pids:\n            if p in tr_map:\n                idx.append(tr_map[p])\n                Y.append(labels[p])\n\n        X = Xtr[idx]\n        Y = np.array(Y)\n\n        Xtrn, Xv, Ytrn, Yv = train_test_split(X, Y, test_size=0.1, random_state=42)\n\n        model = MLP(EMBED, len(terms)).to(DEVICE)\n        opt = optim.Adam(model.parameters(), lr=LR)\n        loss_fn = nn.BCEWithLogitsLoss()\n\n        for _ in range(EPOCHS):\n            for i in range(0, len(Xtrn), BS):\n                xb = torch.tensor(Xtrn[i:i+BS]).to(DEVICE)\n                yb = torch.tensor(Ytrn[i:i+BS]).to(DEVICE)\n                opt.zero_grad()\n                loss = loss_fn(model(xb), yb)\n                loss.backward()\n                opt.step()\n\n        with torch.no_grad():\n            pv = torch.sigmoid(model(torch.tensor(Xv).to(DEVICE))).cpu().numpy()\n        th = best_threshold(Yv, pv)\n\n        with torch.no_grad():\n            preds = []\n            for i in range(0, len(Xte), BS):\n                xb = torch.tensor(Xte[i:i+BS]).to(DEVICE)\n                preds.append(torch.sigmoid(model(xb)).cpu().numpy())\n        preds = np.vstack(preds)\n\n        for i, pid in enumerate(te_ids):\n            idxs = np.where(preds[i] >= th)[0]\n            if len(idxs) == 0:\n                idxs = np.argsort(preds[i])[-2:]\n            for j in idxs:\n                submission.append(f\"{pid}\\t{terms[j]}\\t{preds[i][j]:.3f}\")\n\n    with open(\"submission.tsv\", \"w\") as f:\n        f.write(\"\\n\".join(submission))\n\n    print(\">> submission.tsv created successfully\")\n\nif __name__ == \"__main__\":\n    main()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-23T06:37:47.816866Z","iopub.execute_input":"2025-12-23T06:37:47.817193Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\ndf = pd.read_csv(\n    \"/kaggle/working/submission.tsv\",\n    sep=\"\\t\"\n)\n\nprint(df.columns)\nprint(df.head())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}