{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14875579,"sourceType":"competition"}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q transformers obonet","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"!pip install biopython","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 1: CONFIGURATION & PATHS\n# =========================================================================================\n# Tên block: Config\nimport os\nimport gc\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn as nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast\nfrom transformers import AutoTokenizer, EsmModel\nfrom Bio import SeqIO\nfrom tqdm.auto import tqdm\nfrom sklearn.model_selection import train_test_split\nimport networkx as nx\nimport obonet\nimport logging\n\n# Cấu hình Logging\nlogging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')\nlogger = logging.getLogger()\n\nclass Config:\n    ROOT_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\"\n    \n    TRAIN_FASTA = os.path.join(ROOT_DIR, \"Train/train_sequences.fasta\")\n    TEST_FASTA = os.path.join(ROOT_DIR, \"Test/testsuperset.fasta\")\n    \n    # Các file bổ trợ\n    TRAIN_TERMS = os.path.join(ROOT_DIR, \"Train\", \"train_terms.tsv\")\n    TRAIN_TAXONOMY = os.path.join(ROOT_DIR, \"Train\", \"train_taxonomy.tsv\")\n    GO_OBO = os.path.join(ROOT_DIR, \"Train\", \"go-basic.obo\")\n    IA_FILE = os.path.join(ROOT_DIR, \"IA.tsv\")\n    TEST_TAXONOMY = os.path.join(ROOT_DIR, \"Test\", \"testsuperset-taxon-list.tsv\") # Dùng để lấy ID loài cho tập test\n    \n    SUBMISSION_FILE = \"submission.tsv\"\n    \n    # --- MODEL SETTINGS ---\n    # Sử dụng ESM2-35M (t12) để chạy NHANH trong demo. \n    # ĐỂ ĐẠT ĐIỂM CAO NHẤT: Hãy đổi thành \"facebook/esm2_t30_150M_UR50D\"\n    MODEL_NAME = \"facebook/esm2_t30_150M_UR50D\" \n    \n    MAX_LEN = 1024\n    EMBED_BATCH_SIZE = 32 # Tăng giảm tùy VRAM\n    TRAIN_BATCH_SIZE = 512 # Batch lớn cho training nhanh\n    \n    # --- TRAINING SETTINGS ---\n    NUM_LABELS = 3000   # Số lượng GO Term phổ biến nhất để học\n    EPOCHS = 25\n    LR = 1e-3\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n    \n    # Taxonomy\n    NUM_TAXA = 0 # Sẽ được cập nhật tự động\n\ncfg = Config()\nlogger.info(f\"Running on {cfg.DEVICE} using model {cfg.MODEL_NAME}\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 2: DATA LOADING & PARSING (CUSTOM FOR YOUR FILES)\n# =========================================================================================\n# Tên block: Data Loading\n\ndef load_fasta(path, is_train=False):\n    ids, seqs = [], []\n    # Xử lý path nếu file không nằm ngay root (Kaggle dataset structure varies)\n    if not os.path.exists(path):\n        # Thử tìm trong folder con nếu path sai\n        alt_path = path.replace(cfg.ROOT_DIR, os.path.join(cfg.ROOT_DIR, \"Train\")) if is_train else path.replace(cfg.ROOT_DIR, os.path.join(cfg.ROOT_DIR, \"Test\"))\n        if os.path.exists(alt_path): path = alt_path\n            \n    logger.info(f\"Reading FASTA from: {path}\")\n    for record in SeqIO.parse(path, \"fasta\"):\n        # Format CAFA: \"10116|A0A096MJN4\" -> TaxonID | ProteinID\n        # Nhưng theo ảnh, train_sequences chỉ có ProteinID hoặc format khác\n        # Ta sẽ lấy ID sạch\n        head = record.id.split(\"|\")[1] if \"|\" in record.id else record.id\n        ids.append(str(head))\n        seqs.append(str(record.seq))\n    return ids, seqs\n\ndef load_taxonomy_map(train_tax_path, test_tax_list_path):\n    \"\"\"\n    Tạo map: ProteinID -> TaxonID (Int)\n    \"\"\"\n    logger.info(\"Loading Taxonomy Data...\")\n    \n    # 1. Load Train Taxonomy\n    # Format ảnh: EntryID \\t taxon_id\n    df_train_tax = pd.read_csv(train_tax_path, sep='\\t')\n    # Đảm bảo cột đúng tên (dựa trên ảnh header có vẻ như # 9606 gì đó, cần xử lý header)\n    # Ta sẽ đọc header=0, và rename\n    df_train_tax.columns = ['EntryID', 'taxon_id'] \n    \n    # 2. Xử lý tập Test\n    # File testsuperset-taxon-list.tsv trong ảnh: ID (taxon) \\t Species\n    # Nhưng ta cần biết Protein nào thuộc Taxon nào. \n    # Trong CAFA, header file FASTA test thường chứa TaxonID: \"10116|PROT_ID\"\n    # Nếu file FASTA test không có taxon, ta phải gán mặc định hoặc tìm file ánh xạ khác.\n    # GIẢ ĐỊNH: Header test fasta có format \"TaxonID|ProteinID\". Nếu không, gán taxon phổ biến nhất (9606)\n    \n    # Tạo danh sách tất cả các loài (Train + Test) để encode\n    all_taxa = df_train_tax['taxon_id'].unique().tolist()\n    \n    # Thêm các loài từ file test definition (từ ảnh testsuperset-taxon-list.tsv)\n    df_test_species = pd.read_csv(test_tax_list_path, sep='\\t')\n    # Cột đầu tiên trong ảnh là ID (Taxon ID)\n    test_taxa = df_test_species.iloc[:, 0].unique().tolist()\n    all_taxa.extend(test_taxa)\n    \n    unique_taxa = sorted(list(set(all_taxa)))\n    tax2idx = {t: i for i, t in enumerate(unique_taxa)}\n    \n    return df_train_tax, tax2idx\n\n# --- Load Data ---\ntrain_ids, train_seqs = load_fasta(cfg.TRAIN_FASTA, is_train=True)\ntest_ids, test_seqs = load_fasta(cfg.TEST_FASTA, is_train=False)\n\ndf_train_tax, tax2idx = load_taxonomy_map(cfg.TRAIN_TAXONOMY, cfg.TEST_TAXONOMY)\ncfg.NUM_TAXA = len(tax2idx)\nlogger.info(f\"Total Unique Taxa: {cfg.NUM_TAXA}\")\n\n# Map Train ID -> Taxon Index\ntrain_id_to_tax_idx = {}\n# Tạo dict tra cứu nhanh\ntemp_dict = dict(zip(df_train_tax.EntryID, df_train_tax.taxon_id))\nfor pid in train_ids:\n    # Nếu không tìm thấy, gán 9606 (Human) hoặc 0\n    tid = temp_dict.get(pid, 9606) \n    train_id_to_tax_idx[pid] = tax2idx.get(tid, 0)\n\n# Map Test ID -> Taxon Index (Parse từ Fasta Header hoặc mặc định)\n# Format test fasta CAFA thường là: >TaxID|ProtID\ntest_id_to_tax_idx = {}\nraw_test_headers = [record.id for record in SeqIO.parse(cfg.TEST_FASTA, \"fasta\")] # Đọc lại raw header\nfor raw_h, pid in zip(raw_test_headers, test_ids):\n    if \"|\" in raw_h:\n        try:\n            tid = int(raw_h.split(\"|\")[0])\n            test_id_to_tax_idx[pid] = tax2idx.get(tid, 0)\n        except:\n            test_id_to_tax_idx[pid] = 0 # Fallback\n    else:\n        test_id_to_tax_idx[pid] = 0","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 3: OBO PARSING & GRAPH BUILDING (THE SECRET SAUCE)\n# =========================================================================================\n# Tên block: Ontology Processing\n\nlogger.info(\"Parsing OBO Ontology...\")\ngraph = obonet.read_obo(cfg.GO_OBO)\n\n# Load Train Terms để chọn Top Labels\ndf_terms = pd.read_csv(cfg.TRAIN_TERMS, sep='\\t')\n# Cột: EntryID, term, aspect. Lọc theo tần suất\ntop_terms = df_terms['term'].value_counts().index[:cfg.NUM_LABELS].tolist()\nterm2idx = {t: i for i, t in enumerate(top_terms)}\nidx2term = {i: t for t, i in term2idx.items()}\n\n# Xây dựng ma trận cha-con (Parent-Child Matrix) cho Lan truyền\n# Logic: Nếu A là cha của B, thì Score(A) phải >= Score(B)\n# Ta tạo danh sách: với mỗi term index i, danh sách các cha của nó (mà cũng nằm trong top_terms)\nterm_parents_indices = {}\n\nfor term in tqdm(top_terms, desc=\"Building Ontology Graph\"):\n    if term in graph:\n        # Lấy tất cả tổ tiên (ancestors)\n        ancestors = list(nx.ancestors(graph, term))\n        valid_parents = [term2idx[a] for a in ancestors if a in term2idx]\n        if valid_parents:\n            term_parents_indices[term2idx[term]] = valid_parents\n\nlogger.info(\"Ontology Graph Built.\")\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 4: ESM2 FEATURE EXTRACTION\n# =========================================================================================\n# Tên block: Embedding Extraction\n\ndef extract_esm_embeddings(ids, seqs, save_name):\n    \"\"\"Trích xuất và lưu vào đĩa để tiết kiệm RAM\"\"\"\n    if os.path.exists(save_name):\n        logger.info(f\"Loading cached embeddings: {save_name}\")\n        return np.load(save_name)\n    \n    logger.info(f\"Extracting embeddings for {len(seqs)} sequences...\")\n    tokenizer = AutoTokenizer.from_pretrained(cfg.MODEL_NAME)\n    model = EsmModel.from_pretrained(cfg.MODEL_NAME).to(cfg.DEVICE)\n    model.eval()\n    \n    embeddings = []\n    \n    # Batch processing\n    for i in tqdm(range(0, len(seqs), cfg.EMBED_BATCH_SIZE), desc=\"ESM2 Inference\"):\n        batch_seqs = seqs[i : i + cfg.EMBED_BATCH_SIZE]\n        \n        # Tokenize\n        encoded = tokenizer(\n            batch_seqs, \n            padding=True, \n            truncation=True, \n            max_length=cfg.MAX_LEN, \n            return_tensors='pt'\n        ).to(cfg.DEVICE)\n        \n        with torch.no_grad():\n            with autocast():\n                output = model(**encoded)\n        \n        # Mean Pooling (bỏ qua padding)\n        # output.last_hidden_state: (Batch, Seq, Dim)\n        mask = encoded['attention_mask'].unsqueeze(-1).float()\n        sum_embed = torch.sum(output.last_hidden_state * mask, dim=1)\n        sum_mask = torch.clamp(mask.sum(dim=1), min=1e-9)\n        mean_embed = (sum_embed / sum_mask).cpu().numpy().astype(np.float16)\n        \n        embeddings.append(mean_embed)\n        \n    all_emb = np.concatenate(embeddings, axis=0)\n    np.save(save_name, all_emb)\n    \n    del model, tokenizer\n    torch.cuda.empty_cache()\n    return all_emb\n\nX_train_esm = extract_esm_embeddings(train_ids, train_seqs, \"train_esm.npy\")\nX_test_esm = extract_esm_embeddings(test_ids, test_seqs, \"test_esm.npy\")\n\n# Chuẩn bị input Taxonomy\nX_train_tax = np.array([train_id_to_tax_idx[pid] for pid in train_ids])\nX_test_tax = np.array([test_id_to_tax_idx[pid] for pid in test_ids])\n\n# Chuẩn bị Labels (Multi-hot)\nlogger.info(\"Encoding Labels...\")\nY_train = np.zeros((len(train_ids), cfg.NUM_LABELS), dtype=np.float32)\ntrain_id_map = {pid: i for i, pid in enumerate(train_ids)}\n\n# Filter chỉ lấy rows có trong top terms\ndf_filtered = df_terms[df_terms['term'].isin(top_terms)]\nfor _, row in tqdm(df_filtered.iterrows(), total=len(df_filtered), desc=\"Filling Y_train\"):\n    if row['EntryID'] in train_id_map:\n        r = train_id_map[row['EntryID']]\n        c = term2idx[row['term']]\n        Y_train[r, c] = 1.0\n\n# Split Validation\nX_tr_esm, X_val_esm, X_tr_tax, X_val_tax, y_tr, y_val = train_test_split(\n    X_train_esm, X_train_tax, Y_train, test_size=0.1, random_state=42\n)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 5: TAXONOMY-AWARE RESNET MODEL\n# =========================================================================================\n# Tên block: Model Definition\n\nclass ResidualBlock(nn.Module):\n    def __init__(self, in_features, out_features, dropout=0.3):\n        super().__init__()\n        self.block = nn.Sequential(\n            nn.Linear(in_features, out_features),\n            nn.BatchNorm1d(out_features),\n            nn.ReLU(),\n            nn.Dropout(dropout),\n            nn.Linear(out_features, out_features),\n            nn.BatchNorm1d(out_features),\n            nn.ReLU(),\n            nn.Dropout(dropout)\n        )\n        self.shortcut = nn.Sequential()\n        if in_features != out_features:\n            self.shortcut = nn.Sequential(\n                nn.Linear(in_features, out_features),\n                nn.BatchNorm1d(out_features)\n            )\n\n    def forward(self, x):\n        return self.block(x) + self.shortcut(x)\n\nclass TaxonResNet(nn.Module):\n    def __init__(self, esm_dim, num_taxa, tax_embed_dim, num_labels):\n        super().__init__()\n        # Taxonomy Embedding\n        self.tax_embedding = nn.Embedding(num_taxa, tax_embed_dim)\n        \n        # Combine ESM + Taxon\n        input_dim = esm_dim + tax_embed_dim\n        \n        # ResNet Architecture\n        self.layer1 = ResidualBlock(input_dim, 1024)\n        self.layer2 = ResidualBlock(1024, 512)\n        self.classifier = nn.Linear(512, num_labels)\n        \n    def forward(self, esm, tax):\n        t_emb = self.tax_embedding(tax)\n        x = torch.cat([esm, t_emb], dim=1)\n        x = self.layer1(x)\n        x = self.layer2(x)\n        return torch.sigmoid(self.classifier(x))\n\n# Dataset Wrapper\nclass CAFA_Dataset(Dataset):\n    def __init__(self, esm, tax, labels=None):\n        self.esm = torch.tensor(esm, dtype=torch.float32)\n        self.tax = torch.tensor(tax, dtype=torch.long)\n        self.labels = torch.tensor(labels, dtype=torch.float32) if labels is not None else None\n        \n    def __len__(self): return len(self.esm)\n    \n    def __getitem__(self, idx):\n        if self.labels is not None:\n            return self.esm[idx], self.tax[idx], self.labels[idx]\n        return self.esm[idx], self.tax[idx]","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 6: TRAINING LOOP\n# =========================================================================================\n# Tên block: Training\n\ntrain_ds = CAFA_Dataset(X_tr_esm, X_tr_tax, y_tr)\nval_ds = CAFA_Dataset(X_val_esm, X_val_tax, y_val)\ntrain_loader = DataLoader(train_ds, batch_size=cfg.TRAIN_BATCH_SIZE, shuffle=True)\nval_loader = DataLoader(val_ds, batch_size=cfg.TRAIN_BATCH_SIZE, shuffle=False)\n\nmodel = TaxonResNet(\n    esm_dim=X_train_esm.shape[1], \n    num_taxa=cfg.NUM_TAXA + 1, # +1 cho trường hợp unknown\n    tax_embed_dim=64, \n    num_labels=cfg.NUM_LABELS\n).to(cfg.DEVICE)\n\noptimizer = torch.optim.AdamW(model.parameters(), lr=cfg.LR, weight_decay=1e-4)\nscheduler = torch.optim.lr_scheduler.OneCycleLR(\n    optimizer, max_lr=cfg.LR, steps_per_epoch=len(train_loader), epochs=cfg.EPOCHS\n)\ncriterion = nn.BCELoss()\n\nbest_loss = float('inf')\nlogger.info(\"Starting Training...\")\n\nfor epoch in range(cfg.EPOCHS):\n    model.train()\n    train_loss = 0\n    for esm, tax, lbl in tqdm(train_loader, desc=f\"Epoch {epoch+1}\", leave=False):\n        esm, tax, lbl = esm.to(cfg.DEVICE), tax.to(cfg.DEVICE), lbl.to(cfg.DEVICE)\n        \n        optimizer.zero_grad()\n        preds = model(esm, tax)\n        loss = criterion(preds, lbl)\n        loss.backward()\n        optimizer.step()\n        scheduler.step()\n        train_loss += loss.item()\n        \n    # Validation\n    model.eval()\n    val_loss = 0\n    with torch.no_grad():\n        for esm, tax, lbl in val_loader:\n            esm, tax, lbl = esm.to(cfg.DEVICE), tax.to(cfg.DEVICE), lbl.to(cfg.DEVICE)\n            preds = model(esm, tax)\n            val_loss += criterion(preds, lbl).item()\n            \n    avg_val_loss = val_loss / len(val_loader)\n    logger.info(f\"Epoch {epoch+1}: Train Loss={train_loss/len(train_loader):.4f}, Val Loss={avg_val_loss:.4f}\")\n    \n    if avg_val_loss < best_loss:\n        best_loss = avg_val_loss\n        torch.save(model.state_dict(), \"best_resnet.pth\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# =========================================================================================\n# BLOCK 7: INFERENCE & GRAPH PROPAGATION & SUBMISSION\n# =========================================================================================\n# Tên block: Inference\n\nmodel.load_state_dict(torch.load(\"best_resnet.pth\"))\nmodel.eval()\n\ntest_ds = CAFA_Dataset(X_test_esm, X_test_tax)\ntest_loader = DataLoader(test_ds, batch_size=cfg.TRAIN_BATCH_SIZE * 2, shuffle=False)\n\npreds = []\nlogger.info(\"Running Inference...\")\nwith torch.no_grad():\n    for esm, tax in tqdm(test_loader):\n        esm, tax = esm.to(cfg.DEVICE), tax.to(cfg.DEVICE)\n        out = model(esm, tax)\n        preds.append(out.cpu().numpy())\n\npreds = np.concatenate(preds, axis=0)\n\n# --- Post-processing: Graph Propagation ---\n# Score(Parent) = max(Score(Parent), Score(Child))\n# Ta cần duyệt ngược từ con lên cha. Tuy nhiên, đơn giản hơn là lặp vài lần để lan truyền.\n# Hoặc duyệt theo thứ tự topo. Ở đây dùng cách lặp đơn giản nhưng hiệu quả.\n\nlogger.info(\"Applying Ontology Propagation (This improves F-max)...\")\n# Copy preds để không ghi đè dữ liệu gốc\nfinal_preds = preds.copy()\n\n# Duyệt qua từng protein\nfor i in tqdm(range(final_preds.shape[0]), desc=\"Propagating Scores\"):\n    # Lấy vector dự đoán của 1 protein\n    p_vec = final_preds[i]\n    \n    # Với mỗi term, nếu nó có cha, cập nhật điểm cho cha\n    # Để tối ưu, chỉ xét các term có điểm > 0\n    # Cần lan truyền từ lá lên gốc -> Sắp xếp term theo thứ tự nào đó? \n    # Cách tốt nhất không đệ quy: Loop qua dict term_parents_indices\n    # Tuy nhiên, cần đảm bảo thứ tự. Đơn giản nhất: Loop vài pass.\n    \n    # Pass 1:\n    for child_idx, parent_indices in term_parents_indices.items():\n        child_score = p_vec[child_idx]\n        if child_score > 0.01: # Chỉ xét nếu con có điểm\n            for parent_idx in parent_indices:\n                p_vec[parent_idx] = max(p_vec[parent_idx], child_score)\n    \n    final_preds[i] = p_vec\n\n# --- Write Submission ---\nlogger.info(\"Writing submission.tsv...\")\nwith open(cfg.SUBMISSION_FILE, 'w') as f:\n    for i in tqdm(range(len(test_ids))):\n        # Filter terms > threshold\n        indices = np.where(final_preds[i] >= 0.01)[0]\n        for idx in indices:\n            score = final_preds[i][idx]\n            term = idx2term[idx]\n            f.write(f\"{test_ids[i]}\\t{term}\\t{score:.3f}\\n\")\n\nlogger.info(\"Pipeline Finished! Download submission.tsv\")","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}