{"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"},{"sourceId":13482672,"sourceType":"datasetVersion","datasetId":8559888}],"dockerImageVersionId":31192,"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,"execution":{"iopub.status.busy":"2025-12-12T15:25:18.073201Z","iopub.execute_input":"2025-12-12T15:25:18.073705Z","iopub.status.idle":"2025-12-12T15:25:18.081394Z","shell.execute_reply.started":"2025-12-12T15:25:18.073681Z","shell.execute_reply":"2025-12-12T15:25:18.080794Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def load_fasta_to_df(path):\n    \"\"\"\n    Đọc file FASTA và trả về DataFrame với:\n      - EntryID: mã protein (thường là accession trong header UniProt)\n      - raw_id: toàn bộ header sau dấu '>'\n      - sequence: chuỗi amino acid\n    \"\"\"\n    records = []\n\n    with open(path, \"r\") as f:\n        current_id = None\n        current_seq_parts = []\n\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n\n            if line.startswith(\">\"):\n                # Lưu record cũ nếu có\n                if current_id is not None:\n                    seq = \"\".join(current_seq_parts)\n                    parts = current_id.split(\" \")\n                    entry_id = parts[0] if len(parts) >= 2 else current_id\n\n                    records.append({\n                        \"EntryID\": entry_id,\n                        \"raw_id\": current_id,\n                        \"sequence\": seq\n                    })\n\n                # Bắt đầu record mới\n                current_id = line[1:]  # bỏ dấu '>'\n                current_seq_parts = []\n            else:\n                current_seq_parts.append(line)\n\n        # Lưu record cuối\n        if current_id is not None:\n            seq = \"\".join(current_seq_parts)\n            parts = current_id.split(\" \")\n            entry_id = parts[0] if len(parts) >= 2 else current_id\n\n            records.append({\n                \"EntryID\": entry_id,\n                \"raw_id\": current_id,\n                \"sequence\": seq\n            })\n\n    return pd.DataFrame(records)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:25:19.726437Z","iopub.execute_input":"2025-12-12T15:25:19.727003Z","iopub.status.idle":"2025-12-12T15:25:19.733295Z","shell.execute_reply.started":"2025-12-12T15:25:19.726969Z","shell.execute_reply":"2025-12-12T15:25:19.732601Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_superset = load_fasta_to_df(\"/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta\")\nprint(\"Số sample test :\", len(test_superset))\ntest_superset","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:25:22.153971Z","iopub.execute_input":"2025-12-12T15:25:22.154624Z","iopub.status.idle":"2025-12-12T15:25:24.311363Z","shell.execute_reply.started":"2025-12-12T15:25:22.154601Z","shell.execute_reply":"2025-12-12T15:25:24.310748Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Dùng file obo để tính depth của labels","metadata":{}},{"cell_type":"code","source":"import re\nfrom collections import deque, defaultdict\ndef parse_go_obo(obo_path):\n    \"\"\"\n    Parse file go.obo, trả về:\n      - parents: dict[go_id] = list các is_a parent\n      - namespace: dict[go_id] = 'molecular_function' / 'biological_process' / 'cellular_component'\n    Chỉ dùng is_a (không xử lý part_of để đơn giản).\n    \"\"\"\n    parents = {}\n    namespace = {}\n    current_id = None\n\n    with open(obo_path, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if line == \"[Term]\":\n                current_id = None\n                continue\n            if not line or line.startswith(\"!\"):\n                continue\n\n            if line.startswith(\"id: GO:\"):\n                current_id = line.split(\"id: \")[1]\n                parents.setdefault(current_id, [])\n            elif line.startswith(\"namespace:\") and current_id is not None:\n                ns = line.split(\"namespace: \")[1].strip()\n                namespace[current_id] = ns\n            elif line.startswith(\"is_a:\") and current_id is not None:\n                # dòng dạng \"is_a: GO:0008150 ! biological_process\"\n                m = re.match(r\"is_a:\\s+(GO:\\d+)\", line)\n                if m:\n                    p = m.group(1)\n                    parents[current_id].append(p)\n\n    return parents, namespace\nobo_path = \"/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo\"\nparents, namespaces = parse_go_obo(obo_path)\nprint(\"Số GO terms:\", len(parents))\n\ndef compute_go_depths(parents, namespaces):\n    \"\"\"\n    Tính depth cho từng GO dựa trên is_a và root:\n      - GO:0003674 (MFO)\n      - GO:0008150 (BPO)\n      - GO:0005575 (CCO)\n\n    Trả về:\n      - depth: dict[go_id] = depth (int, 0 cho root, 1 cho con trực tiếp, ...)\n    \"\"\"\n\n    roots = {\n        \"GO:0003674\",  # molecular_function\n        \"GO:0008150\",  # biological_process\n        \"GO:0005575\",  # cellular_component\n    }\n\n    # build children graph (ngược với parents)\n    children = defaultdict(list)\n    for child, ps in parents.items():\n        for p in ps:\n            children[p].append(child)\n\n    depth = {}\n    q = deque()\n\n    # khởi tạo depth = 0 cho các root\n    for r in roots:\n        if r in parents:\n            depth[r] = 0\n            q.append(r)\n\n    # BFS xuống dưới\n    while q:\n        node = q.popleft()\n        d = depth[node]\n        for c in children.get(node, []):\n            # nếu đã có depth thì lấy min (đường ngắn nhất)\n            if c not in depth or d + 1 < depth[c]:\n                depth[c] = d + 1\n                q.append(c)\n\n    return depth\ndepths = compute_go_depths(parents, namespaces)\nprint(\"Ví dụ depth:\", list(depths.items())[:10])\ndepth_df = pd.DataFrame(\n    [(go, d) for go, d in depths.items()],\n    columns=[\"term\", \"depth\"]\n)\ndepth_df","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:07.133250Z","iopub.execute_input":"2025-12-12T15:33:07.133958Z","iopub.status.idle":"2025-12-12T15:33:07.536511Z","shell.execute_reply.started":"2025-12-12T15:33:07.133933Z","shell.execute_reply":"2025-12-12T15:33:07.535916Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_terms = pd.read_csv(\n    \"/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv\",\n    sep=\"\\t\"\n)\nprint(train_terms.shape)\nprint(len(train_terms[\"term\"].unique()))\nweights = pd.read_csv(\"/kaggle/input/cafa-6-protein-function-prediction/IA.tsv\", sep=\"\\t\")\nweights = weights.rename(columns={\"GO:0000001\": \"term\", \"0.0\": \"weight\"})\nweights = weights.merge(depth_df,on=\"term\", how=\"left\")\nweights.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:10.549968Z","iopub.execute_input":"2025-12-12T15:33:10.550495Z","iopub.status.idle":"2025-12-12T15:33:10.866237Z","shell.execute_reply.started":"2025-12-12T15:33:10.550473Z","shell.execute_reply":"2025-12-12T15:33:10.865597Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Lọc nhãn","metadata":{}},{"cell_type":"code","source":"def select_labels_by_freq_and_weight(\n    train_terms: pd.DataFrame,\n    weights: pd.DataFrame,\n    freq_min: int = 10,\n    weight_min: float = 5.0,\n    weight_max: float = 100.0,\n    min_depth: int = 3,\n):\n    # Bỏ cột EntryID, chỉ giữ term để đếm\n    df_new = train_terms.drop(columns=[\"EntryID\"])\n    \n    # Đếm tần số xuất hiện của từng term\n    term_counts = df_new[\"term\"].value_counts().reset_index()\n    term_counts.columns = [\"term\", \"frequency\"]\n    print(\"📌 term_counts:\", term_counts.shape)\n    print(term_counts.head())\n\n    # Gộp với weight\n    term_counts_weight = (\n        term_counts\n        .merge(weights, on=\"term\", how=\"left\")\n        .fillna(0)\n        .sort_values(\"weight\", ascending=False)\n    )\n    print(\"📌 term_counts_weight:\", term_counts_weight.shape)\n    print(term_counts_weight.head())\n\n    # Lọc theo frequency và weight\n    filtered = term_counts_weight[\n        (term_counts_weight[\"frequency\"] >= freq_min) &\n        (term_counts_weight[\"weight\"] >= weight_min) &\n        (term_counts_weight[\"weight\"] < weight_max) &\n        (term_counts_weight[\"depth\"] >= min_depth)\n    ].reset_index(drop=True)\n\n    selected_labels = filtered[\"term\"].tolist()\n    print(f\"✅ Số nhãn được chọn: {len(selected_labels)} (freq ≥ {freq_min}, weight ≥ {weight_min}, weight <= {weight_max}, depth >= {min_depth})\")\n\n    return selected_labels\n\nselected_labels_batch1 = select_labels_by_freq_and_weight(train_terms, weights, 10, 3, 100, 4)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:12.906051Z","iopub.execute_input":"2025-12-12T15:33:12.906338Z","iopub.status.idle":"2025-12-12T15:33:13.022025Z","shell.execute_reply.started":"2025-12-12T15:33:12.906318Z","shell.execute_reply":"2025-12-12T15:33:13.021416Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_labels = (\n    train_terms\n    .groupby([\"EntryID\",\"aspect\"])[\"term\"]\n    .apply(list)\n    .reset_index()\n)\nprint(train_labels.shape)\ntrain_labels","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:15.393379Z","iopub.execute_input":"2025-12-12T15:33:15.393917Z","iopub.status.idle":"2025-12-12T15:33:18.535349Z","shell.execute_reply.started":"2025-12-12T15:33:15.393896Z","shell.execute_reply":"2025-12-12T15:33:18.534748Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"def train_labels_filtered(\n    train_labels: pd.DataFrame,\n    selected_labels: list\n):\n    train_labels_filtered = train_labels.copy()\n    \n    train_labels_filtered[\"term\"] = train_labels_filtered[\"term\"].apply(\n        lambda terms: [t for t in terms if t in selected_labels]\n    )\n    train_labels_filtered = train_labels_filtered[\n        train_labels_filtered[\"term\"].map(len) > 0\n    ].reset_index(drop=True)\n    print(train_labels_filtered.shape)\n    return train_labels_filtered\n\ntrain_labels_filtered_b1 = train_labels_filtered(train_labels, selected_labels_batch1)\ntrain_labels_filtered_b1.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:18.536475Z","iopub.execute_input":"2025-12-12T15:33:18.537100Z","iopub.status.idle":"2025-12-12T15:33:27.063946Z","shell.execute_reply.started":"2025-12-12T15:33:18.537080Z","shell.execute_reply":"2025-12-12T15:33:27.063335Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np\nimport pandas as pd\nfrom pathlib import Path\n\ndata_dir = Path('/kaggle/input/cafa6-protein-embeddings-esm2')\n\nprotein_ids = pd.read_csv(data_dir / \"protein_ids.csv\")[\"protein_id\"].tolist()\nembeddings = np.load(data_dir / \"protein_embeddings.npy\")\n\n\nprint(f\"Loaded {len(protein_ids)} embeddings of dimension {embeddings.shape[1]}\")\n\npid_emb_dict = {}\nfor pid, emb in zip(protein_ids, embeddings):\n    pid_emb_dict[str(pid)] = emb","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:27.065067Z","iopub.execute_input":"2025-12-12T15:33:27.065288Z","iopub.status.idle":"2025-12-12T15:33:36.150909Z","shell.execute_reply.started":"2025-12-12T15:33:27.065272Z","shell.execute_reply":"2025-12-12T15:33:36.150179Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"emb_df = (\n    pd.DataFrame(\n        [(k, v) for k,v in pid_emb_dict.items()],\n        columns=[\"EntryID\", \"Emb\"]\n    )\n)\nemb_df.head()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:36.151607Z","iopub.execute_input":"2025-12-12T15:33:36.151893Z","iopub.status.idle":"2025-12-12T15:33:36.242728Z","shell.execute_reply.started":"2025-12-12T15:33:36.151874Z","shell.execute_reply":"2025-12-12T15:33:36.242158Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"train_labels_filtered_b1 = train_labels_filtered_b1.merge(emb_df, on=\"EntryID\", how=\"left\")\ntrain_labels_filtered_b1","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:40.426079Z","iopub.execute_input":"2025-12-12T15:33:40.426351Z","iopub.status.idle":"2025-12-12T15:33:40.536725Z","shell.execute_reply.started":"2025-12-12T15:33:40.426330Z","shell.execute_reply":"2025-12-12T15:33:40.536199Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# Kiểm tra phiên bản PyTorch và CUDA hiện tại\nimport torch\n\nTORCH_VERSION = torch.__version__\nCUDA_VERSION = torch.version.cuda\nprint(f\"PyTorch Version: {TORCH_VERSION}, CUDA Version: {CUDA_VERSION}\")\n\n# Cài đặt các thư viện phụ thuộc (PyTorch Scatter, PyTorch Sparse)\n# LƯU Ý: Thay thế {TORCH_VERSION} bằng phiên bản thực tế của bạn (ví dụ: 2.1.0)\n# và {CUDA_VERSION} bằng phiên bản thực tế của bạn (ví dụ: 12.1)\n\n# Ví dụ cho PyTorch 2.1.0 và CUDA 12.1 (Hãy thay thế bằng phiên bản của bạn!)\n# Tùy thuộc vào phiên bản, bạn có thể cần tìm lại lệnh cài đặt chính xác trên trang chủ PyG\n# Đây là cú pháp phổ biến trên nhiều nền tảng:\n\n# Cài đặt PyTorch Scatter\n!pip install torch_scatter -f https://data.pyg.org/whl/torch-{TORCH_VERSION}.html\n# Cài đặt PyTorch Sparse\n!pip install torch_sparse -f https://data.pyg.org/whl/torch-{TORCH_VERSION}.html\n# Cài đặt PyTorch Geometric\n!pip install torch_geometric obonet networkx\n\nprint(\"\\nĐã cố gắng cài đặt torch_geometric và các thư viện liên quan.\")\n\n# Khởi động lại runtime/kernel sau khi cài đặt (Khuyến nghị)\n# Bạn nên chạy lệnh này, sau đó dùng tùy chọn \"Restart Kernel\" của Notebook\n# import os\n# os._exit(0)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:42.858341Z","iopub.execute_input":"2025-12-12T15:33:42.858962Z","iopub.status.idle":"2025-12-12T15:33:55.865326Z","shell.execute_reply.started":"2025-12-12T15:33:42.858940Z","shell.execute_reply":"2025-12-12T15:33:55.864564Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch_geometric.nn import GCNConv\nimport torch.nn.functional as F\nimport torch_geometric\nfrom typing import Dict, List, Tuple\nfrom sklearn.model_selection import train_test_split\nimport obonet # Cần cho bước 5\n\n# --- THIẾT LẬP CƠ BẢN ---\nESM2_DIM = 1280 \nBATCH_SIZE = 64\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nOBO_FILE_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'\n\n# GÁN DATAFRAME CỦA BẠN VÀO BIẾN `df`\ndf = train_labels_filtered_b1.copy() \n\n# Khởi tạo một dictionary để chứa tất cả dữ liệu đã chia theo Aspect\n# --- LOGIC CODE ĐIỀU CHỈNH CỤ THỂ ---\n\n# Khởi tạo dictionaries để chứa data cho từng Aspect\naspect_data = {'C': {}, 'P': {}, 'F': {}}\naspects_list = ['C', 'P', 'F']\n\n# Lấy danh sách EntryID duy nhất và Embeddings tương ứng (Đây là X chung)\n# Vì EntryID lặp lại, ta chỉ lấy 1 lần.\nX = np.stack(df['Emb'].values)\nall_entry_ids = df['EntryID'].values\n\n# Lấy danh sách các EntryID duy nhất và Embeddings tương ứng (Tạo X chung)\nunique_entries_df = df.drop_duplicates(subset=['EntryID']).set_index('EntryID')\nX_unique = np.stack(unique_entries_df['Emb'].values)\nunique_entry_ids = unique_entries_df.index.tolist()\ndata_size_unique = len(unique_entry_ids)\n# Ánh xạ EntryID -> Index để xây dựng ma trận Y\nentry_id_to_index = {eid: i for i, eid in enumerate(unique_entry_ids)}\n\n\n# 1. TẠO LABEL ENCODER & Y_MATRIX RIÊNG BIỆT CHO MỖI ASPECT\nprint(\"1. Đang tạo Label Encoder và Y_Matrix riêng biệt cho C, P, F...\")\n\nfor aspect in aspects_list:\n    # Lọc DataFrame chỉ giữ lại các hàng của Aspect này\n    df_aspect = df[df['aspect'] == aspect].copy()\n    \n    # Lấy TẤT CẢ GO Terms duy nhất cho Aspect này\n    go_terms_all = sorted(list(set([term for sublist in df_aspect['term'] for term in sublist])))\n    go_id_to_index = {go_id: i for i, go_id in enumerate(go_terms_all)}\n    \n    num_labels = len(go_terms_all)\n    \n    aspect_data[aspect]['go_id_to_index'] = go_id_to_index\n    aspect_data[aspect]['num_labels'] = num_labels\n    aspect_data[aspect]['go_terms_all'] = go_terms_all\n    \n    # 2. Xây dựng Ma trận Nhãn Y_aspect (Kích thước: Unique_Proteins x Num_Labels)\n    Y_aspect = np.zeros((data_size_unique, num_labels), dtype=np.float32)\n\n    for i, row in df_aspect.iterrows():\n        entry_idx = entry_id_to_index.get(row['EntryID'])\n        if entry_idx is not None:\n            for go_id in row['term']:\n                if go_id in go_id_to_index:\n                    Y_aspect[entry_idx, go_id_to_index[go_id]] = 1\n                    \n    aspect_data[aspect]['Y'] = Y_aspect\n    print(f\" -> Aspect {aspect}: {num_labels} nhãn. Y_Matrix size: {Y_aspect.shape}\")\n\n\n# 3. CHIA DỮ LIỆU CHUNG VÀ TÍNH WEIGHTS\nprint(\"3. Đang chia dữ liệu và tính toán Weights...\")\n# Chia Index (và EntryID) để đảm bảo X_train/X_val là nhất quán\ntrain_indices, val_indices = train_test_split(\n    np.arange(data_size_unique), test_size=0.2, random_state=42\n)\n\nX_train = X_unique[train_indices]\nX_val = X_unique[val_indices]\n\nfor aspect in aspects_list:\n    Y_aspect = aspect_data[aspect]['Y']\n    \n    # Chia Y_aspect theo các index đã chia ở trên\n    Y_train_aspect = Y_aspect[train_indices]\n    Y_val_aspect = Y_aspect[val_indices]\n    \n    aspect_data[aspect]['Y_train'] = Y_train_aspect\n    aspect_data[aspect]['Y_val'] = Y_val_aspect\n    \n    # Tính Trọng số (Weights)\n    label_counts = Y_train_aspect.sum(axis=0)\n    weights_np = len(Y_train_aspect) / (label_counts + 1e-6)\n    weights_tensor = torch.tensor(weights_np, dtype=torch.float32).to(DEVICE)\n    aspect_data[aspect]['weights_tensor'] = weights_tensor\n    \n    print(f\" -> Aspect {aspect} | Train Samples: {len(Y_train_aspect)}\")\n\n# [Sau đó, bạn tiếp tục với các bước 4, 5, 6, 7, 8 (GoGCN Model) như đã viết]\n\n\n# 4. PYTORCH DATALOADER CHUNG\nclass CAFADatasetMulti(Dataset):\n    def __init__(self, features, labels_c, labels_f, labels_p):\n        self.features = torch.tensor(features, dtype=torch.float32)\n        self.labels_c = torch.tensor(labels_c, dtype=torch.float32)\n        self.labels_f = torch.tensor(labels_f, dtype=torch.float32)\n        self.labels_p = torch.tensor(labels_p, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        return (self.features[idx], self.labels_c[idx], self.labels_f[idx], self.labels_p[idx])\n\ntrain_dataset = CAFADatasetMulti(\n    X_train, \n    aspect_data['C']['Y_train'], \n    aspect_data['F']['Y_train'], \n    aspect_data['P']['Y_train']\n)\ntrain_loader = DataLoader(train_dataset, batch_size=BATCH_SIZE, shuffle=True)\n\nval_dataset = CAFADatasetMulti(\n    X_val, \n    aspect_data['C']['Y_val'], \n    aspect_data['F']['Y_val'], \n    aspect_data['P']['Y_val']\n)\nval_loader = DataLoader(val_dataset, batch_size=BATCH_SIZE, shuffle=False)\n\n# 5. XÂY DỰNG EDGE INDEX RIÊNG BIỆT CHO MỖI ASPECT\nprint(\"5. Đang xây dựng Edge Index riêng cho từng Aspect...\")\n\n# --- THIẾT LẬP ĐƯỜNG DẪN OBO (BẮT BUỘC PHẢI CHỈNH SỬA) ---\n# Vui lòng thay thế đường dẫn này bằng đường dẫn thực tế của file OBO trên Kaggle của bạn.\nOBO_FILE_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo' \n\ndef load_go_graph_obo(obo_file_path: str):\n    \"\"\"\n    Tải đồ thị GO gốc (Multigraph) từ file OBO bằng obonet.\n    \"\"\"\n    print(f\"-> Đang cố gắng tải đồ thị GO từ: {obo_file_path}\")\n    try:\n        # obonet.read_obo() trả về một đối tượng networkx.MultiDiGraph\n        go_graph = obonet.read_obo(obo_file_path)\n        print(f\"-> Tải thành công. Đồ thị có {go_graph.number_of_nodes()} GO Terms.\")\n        return go_graph\n    except FileNotFoundError:\n        print(f\"LỖI: Không tìm thấy file OBO tại '{obo_file_path}'. Quá trình tạo Edge Index sẽ bị bỏ qua.\")\n        return None\n    except Exception as e:\n        print(f\"LỖI KHÁC khi tải file OBO: {e}\")\n        return None\n\n# --- THỰC THI: TẢI ĐỒ THỊ VÀ GÁN VÀO BIẾN go_graph ---\ngo_graph = load_go_graph_obo(OBO_FILE_PATH)\n\n# Kiểm tra nếu go_graph tải thành công để quyết định logic tiếp theo\nif go_graph is None:\n    # Nếu không tải được, go_graph = None. Logic của bạn sẽ dùng Edge Index giả lập.\n    print(\"CẢNH BÁO: Đồ thị GO không được tải. Sử dụng Edge Index ngẫu nhiên.\")\nelse:\n    print(\"Đồ thị GO đã được tải thành công cho việc xây dựng Edge Index theo Aspect.\")\n\n# --- Đảm bảo các thư viện cần thiết đã được import ---\n# import torch_geometric \n# ...\ndef build_go_graph_edges_aspect(go_id_to_index, go_graph, aspect):\n    # Sử dụng logic OBOnet đã sửa lỗi (cần obonet và go_graph đã load)\n    # LỌC ĐỒ THỊ CHỈ CHO ASPECT NÀY\n    \n    num_nodes = len(go_id_to_index)\n    edges_list = []\n    \n    for source_id, target_id, relation in go_graph.edges(keys=True):\n        if relation in ['is_a', 'part_of']:\n            # Chỉ xét các GO ID có trong tập nhãn đã lọc\n            if source_id in go_id_to_index and target_id in go_id_to_index:\n                source_idx = go_id_to_index[source_id]\n                target_idx = go_id_to_index[target_id]\n                edges_list.append((source_idx, target_idx))\n                edges_list.append((target_idx, source_idx)) # Vô hướng cho GCN\n    \n    if not edges_list:\n        # Thêm self-loops nếu không có cạnh nào được tìm thấy\n        edges_list = [(i, i) for i in range(num_nodes)]\n\n    edge_index = torch.tensor(list(set(edges_list)), dtype=torch.long).t().contiguous()\n    # Thêm self-loops và loại bỏ trùng lặp (chắc chắn)\n    edge_index = torch_geometric.utils.coalesce(edge_index, num_nodes=num_nodes)\n    \n    print(f\" -> Aspect {aspect} Edge Count: {edge_index.size(1)}\")\n    return edge_index.to(DEVICE)\n\n# Tạo các Edge Index\nif go_graph:\n    edge_indices = {\n        aspect: build_go_graph_edges_aspect(aspect_data[aspect]['go_id_to_index'], go_graph, aspect)\n        for aspect in aspects_list\n    }\nelse:\n    # Nếu không load được OBO, tạo Edge Index giả lập để code không lỗi\n    print(\"WARNING: Không thể tạo Edge Index thực tế. Tạo Edge Index ngẫu nhiên.\")\n    edge_indices = {\n        aspect: torch_geometric.utils.coalesce(torch.randint(0, aspect_data[aspect]['num_labels'], (2, 2 * aspect_data[aspect]['num_labels'])).long(), num_nodes=aspect_data[aspect]['num_labels']).to(DEVICE)\n        for aspect in aspects_list\n    }\n# Cần đảm bảo GCNConv, torch.nn.functional đã được import\n\nclass MultiHeadGoGCN(torch.nn.Module):\n    def __init__(self, in_dim: int, num_labels: Dict[str, int], edge_indices: Dict[str, torch.Tensor]):\n        super(MultiHeadGoGCN, self).__init__()\n        \n        # 1. Base Feature Extractor (Lớp chung)\n        # Học các đặc trưng chung từ ESM embedding (Kích thước: 1280 -> 1024)\n        self.base_extractor = torch.nn.Sequential(\n            torch.nn.Linear(in_dim, 1024),\n            torch.nn.BatchNorm1d(1024),\n            torch.nn.ReLU(),\n            torch.nn.Dropout(0.4),\n        )\n        \n        # 2. Các Head và GCN độc lập\n        self.gcn_c = self._create_aspect_head(1024, num_labels['C'], edge_indices['C'], 'C')\n        self.gcn_f = self._create_aspect_head(1024, num_labels['F'], edge_indices['F'], 'F')\n        self.gcn_p = self._create_aspect_head(1024, num_labels['P'], edge_indices['P'], 'P')\n\n    def _create_aspect_head(self, in_dim: int, out_dim: int, edge_index_go: torch.Tensor, aspect: str):\n        # Đảm bảo edge_index và V_init được đăng ký là buffer/parameter\n        self.register_buffer(f'edge_index_go_{aspect}', edge_index_go)\n        V_init = torch.nn.Parameter(torch.eye(out_dim), requires_grad=False)\n        self.register_buffer(f'V_init_{aspect}', V_init)\n        \n        return torch.nn.ModuleDict({\n            'head': torch.nn.Sequential(\n                torch.nn.Linear(in_dim, 512),\n                torch.nn.BatchNorm1d(512),\n                torch.nn.ReLU(),\n                torch.nn.Dropout(0.3),\n                torch.nn.Linear(512, out_dim)\n            ),\n            'conv1': GCNConv(out_dim, 512),\n            'conv2': GCNConv(512, 256),\n            'conv3': GCNConv(256, out_dim),\n        })\n\n    def _run_gcn_head(self, x: torch.Tensor, head_module: torch.nn.ModuleDict, aspect: str):\n        \n        edge_index = getattr(self, f'edge_index_go_{aspect}')\n        V_init = getattr(self, f'V_init_{aspect}')\n        \n        # 1. Initial Classification Head (h0)\n        h0 = head_module['head'](x) # (Batch_size, num_labels_aspect)\n        \n        # 2. Học Ma trận Trọng số GO (V_final) bằng GCN\n        V = V_init\n        \n        V = F.relu(head_module['conv1'](V, edge_index))\n        V = F.dropout(V, p=0.3, training=self.training)\n        \n        V = F.relu(head_module['conv2'](V, edge_index))\n        V = F.dropout(V, p=0.3, training=self.training)\n        \n        V_final = head_module['conv3'](V, edge_index) # (num_labels_aspect, num_labels_aspect)\n        \n        # 3. Kết hợp\n        h_final = torch.matmul(h0, V_final)\n        \n        return h_final # Logits\n\n    def forward(self, x_esm: torch.Tensor):\n        # Base features\n        base_features = self.base_extractor(x_esm)\n        \n        # Output cho từng Head\n        out_C = self._run_gcn_head(base_features, self.gcn_c, 'C')\n        out_F = self._run_gcn_head(base_features, self.gcn_f, 'F')\n        out_P = self._run_gcn_head(base_features, self.gcn_p, 'P')\n        \n        return out_C, out_F, out_P\n\n# 2. Khởi tạo Mô hình và Loss Functions\nnum_labels_dict = {a: aspect_data[a]['num_labels'] for a in aspects_list}\n\nmodel = MultiHeadGoGCN(\n    ESM2_DIM, \n    num_labels_dict, \n    edge_indices\n).to(DEVICE)\n\n# Khởi tạo Optimizer (Thêm L2 Regularization)\noptimizer = torch.optim.Adam(model.parameters(), lr=1e-3, weight_decay=1e-5)\n\n# Khởi tạo Loss Functions riêng biệt (Sử dụng IA Weights từ aspect_data)\ncriterion_C = torch.nn.BCEWithLogitsLoss(weight=aspect_data['C']['weights_tensor'])\ncriterion_F = torch.nn.BCEWithLogitsLoss(weight=aspect_data['F']['weights_tensor'])\ncriterion_P = torch.nn.BCEWithLogitsLoss(weight=aspect_data['P']['weights_tensor'])\n\nprint(f\"\\n✅ Mô hình Multi-Head GoGCN đã sẵn sàng ({sum(num_labels_dict.values())} nhãn tổng cộng).\")\n\nNUM_EPOCHS = 100\n# 3. Vòng lặp Huấn luyện (Training Loop)\nprint(f\"Bắt đầu huấn luyện {NUM_EPOCHS} epochs trên {DEVICE}...\")\n\ndef train_epoch(model, loader, criteria, optimizer, device):\n    model.train()\n    total_loss = 0\n    total_loss_c, total_loss_f, total_loss_p = 0, 0, 0\n    \n    for features, labels_c, labels_f, labels_p in loader:\n        \n        # Chuyển dữ liệu lên GPU\n        features = features.to(device)\n        labels_c, labels_f, labels_p = labels_c.to(device), labels_f.to(device), labels_p.to(device)\n        \n        optimizer.zero_grad()\n        \n        # Forward pass\n        output_C, output_F, output_P = model(features)\n        \n        # Tính Loss riêng biệt cho từng Aspect\n        loss_C = criteria['C'](output_C, labels_c)\n        loss_F = criteria['F'](output_F, labels_f)\n        loss_P = criteria['P'](output_P, labels_p)\n        \n        # Tổng Loss\n        total_loss_batch = loss_C + loss_F + loss_P \n        \n        # Backward pass\n        total_loss_batch.backward()\n        optimizer.step()\n        \n        total_loss += total_loss_batch.item()\n        total_loss_c += loss_C.item()\n        total_loss_f += loss_F.item()\n        total_loss_p += loss_P.item()\n\n    avg_loss = total_loss / len(loader)\n    avg_loss_c = total_loss_c / len(loader)\n    avg_loss_f = total_loss_f / len(loader)\n    avg_loss_p = total_loss_p / len(loader)\n    \n    return avg_loss, avg_loss_c, avg_loss_f, avg_loss_p\n\ncriteria = {'C': criterion_C, 'F': criterion_F, 'P': criterion_P}\n\n# --- Vòng lặp Chính ---\nfor epoch in range(1, NUM_EPOCHS + 1):\n    avg_loss, loss_c, loss_f, loss_p = train_epoch(model, train_loader, criteria, optimizer, DEVICE)\n    \n    print(f'Epoch {epoch:03d} | Total Loss: {avg_loss:.4f} | C: {loss_c:.4f}, F: {loss_f:.4f}, P: {loss_p:.4f}')\n\nprint(\"\\nQuá trình huấn luyện Multi-Head GoGCN đã hoàn tất.\")\n\n# --- Định nghĩa mô hình Multi-Head GoGCN và vòng lặp huấn luyện sẽ được thực hiện sau ---\n# (Sử dụng kiến trúc MultiHeadGoGCN đã đề xuất trước đó)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:33:55.866846Z","iopub.execute_input":"2025-12-12T15:33:55.867127Z","iopub.status.idle":"2025-12-12T15:40:40.969463Z","shell.execute_reply.started":"2025-12-12T15:33:55.867095Z","shell.execute_reply":"2025-12-12T15:40:40.968783Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"test_df = pd.DataFrame(test_superset[\"EntryID\"], columns=[\"EntryID\"]).merge(emb_df, on=\"EntryID\", how =\"inner\")\ntest_df ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:42:32.863863Z","iopub.execute_input":"2025-12-12T15:42:32.864812Z","iopub.status.idle":"2025-12-12T15:42:33.116663Z","shell.execute_reply.started":"2025-12-12T15:42:32.864786Z","shell.execute_reply":"2025-12-12T15:42:33.116058Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nimport torch\nfrom torch.utils.data import Dataset, DataLoader\nimport networkx as nx\nimport obonet\nfrom typing import Dict, List, Tuple\nimport torch.nn.functional as F\n\n# --- CÁC THIẾT LẬP VÀ BIẾN CƠ BẢN (Đã định nghĩa trước đó) ---\nBATCH_SIZE = 64\nDEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nOBO_FILE_PATH = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo' # VUI LÒNG ĐIỀU CHỈNH ĐƯỜNG DẪN NÀY\n\n\n# --- 1. HÀM TẠO ĐỒ THỊ GO CÓ HƯỚNG (CHILD -> PARENT) ---\ndef load_go_graph_directed(obo_file_path: str) -> nx.DiGraph:\n    \"\"\"\n    Tải đồ thị GO từ OBO và tạo đồ thị có hướng (Child -> Parent)\n    chỉ sử dụng quan hệ 'is_a' và 'part_of' (để lan truyền).\n    \"\"\"\n    print(f\"1. Đang tải đồ thị GO từ: {obo_file_path}\")\n    try:\n        # obonet tải đồ thị gốc\n        go_graph_multi = obonet.read_obo(obo_file_path)\n    except FileNotFoundError:\n        print(\"LỖI: Không tìm thấy file OBO. Không thể tạo đồ thị lan truyền.\")\n        return None\n\n    # Tạo đồ thị có hướng đơn giản\n    go_graph_directed = nx.DiGraph()\n\n    for source, target, relation in go_graph_multi.edges(keys=True):\n        \n        # Chỉ quan tâm đến các quan hệ hierarchy\n        if relation in ['is_a', 'part_of']:\n            # obonet tải theo chiều Parent -> Child. Ta cần đảo ngược: Child -> Parent\n            go_graph_directed.add_edge(source, target)\n\n    print(f\"2. Đồ thị Lan truyền đã tạo: {go_graph_directed.number_of_nodes()} nodes, {go_graph_directed.number_of_edges()} edges.\")\n    return go_graph_directed\n\n# --- THỰC THI: KHỞI TẠO BIẾN go_graph_directed ---\ngo_graph_directed = load_go_graph_directed(OBO_FILE_PATH)\n\nif go_graph_directed is None:\n    # Nếu lỗi, đặt một đối tượng rỗng để code không bị lỗi NameError\n    go_graph_directed = nx.DiGraph()\n    print(\"CẢNH BÁO: Sử dụng đồ thị rỗng. Quá trình lan truyền sẽ bị bỏ qua hoặc không chính xác.\")\n\n\n# --- 2. HÀM LAN TRUYỀN XÁC SUẤT (PROPAGATE SCORES) ---\ndef propagate_scores_to_parent(scores_matrix: np.ndarray, go_graph_directed: nx.DiGraph, go_id_to_index: dict, go_terms_all: list) -> np.ndarray:\n    \"\"\"\n    Lan truyền score từ con lên cha trong ma trận dự đoán bằng cách duyệt tô pô.\n    scores_matrix: (Số proteins x Số nhãn)\n    \"\"\"\n    \n    if go_graph_directed.number_of_nodes() == 0:\n        print(\"BỎ QUA: Đồ thị rỗng, không thể lan truyền.\")\n        return scores_matrix\n    \n    N_proteins, N_labels = scores_matrix.shape\n    propagated_scores = scores_matrix.copy()\n    \n    try:\n        # 1. Lấy thứ tự tô pô (để đi từ lá lên gốc)\n        topo_order = list(nx.topological_sort(go_graph_directed))\n    except nx.NetworkXUnfeasible:\n        print(\"CẢNH BÁO: Đồ thị GO có chu trình. Bỏ qua lan truyền.\")\n        return scores_matrix \n\n    topo_order.reverse() # Đảo ngược thứ tự: Lá -> Gốc\n    \n    # 2. Lan truyền score\n    for go_id_child in topo_order:\n        \n        # Chỉ lan truyền các GO term có trong tập nhãn đã lọc của bạn\n        if go_id_child not in go_id_to_index:\n            continue\n            \n        idx_child = go_id_to_index[go_id_child]\n        \n        # neighbors(child) là các nút cha (Parent) trong đồ thị Child->Parent\n        for go_id_parent in go_graph_directed.neighbors(go_id_child):\n            \n            # Chỉ lan truyền lên các nút cha có trong tập nhãn đã lọc\n            if go_id_parent not in go_id_to_index:\n                continue\n            \n            idx_parent = go_id_to_index[go_id_parent]\n            \n            # P(A) = max(P(A_old), P(C))\n            propagated_scores[:, idx_parent] = np.maximum(\n                propagated_scores[:, idx_parent],\n                propagated_scores[:, idx_child]\n            )\n            \n    print(\"Lan truyền xác suất đã hoàn tất.\")\n    return propagated_scores\n\n# *******************************************************************\n# BÂY GIỜ, BẠN CÓ THỂ TIẾP TỤC VỚI BƯỚC DỰ ĐOÁN CUỐI CÙNG\n# (Sử dụng hàm này để gọi lại quá trình tạo submission)\n# *******************************************************************","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:42:35.638723Z","iopub.execute_input":"2025-12-12T15:42:35.639244Z","iopub.status.idle":"2025-12-12T15:42:40.720389Z","shell.execute_reply.started":"2025-12-12T15:42:35.639219Z","shell.execute_reply":"2025-12-12T15:42:40.719569Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"Loc theo function","metadata":{}},{"cell_type":"code","source":"class CAFADatasetTest(Dataset):\n    \"\"\"Dataset cho tập Test (chỉ có features)\"\"\"\n    def __init__(self, features):\n        self.features = torch.tensor(features, dtype=torch.float32)\n\n    def __len__(self):\n        return len(self.features)\n\n    def __getitem__(self, idx):\n        # Trả về dummy labels (tensor 0) để phù hợp với cấu trúc DataLoader\n        return self.features[idx], torch.zeros(1) \n\ndef get_test_predictions(model, test_loader, device):\n    \"\"\"Lấy 3 ma trận Logits (C, F, P) từ mô hình Multi-Head trên tập Test.\"\"\"\n    model.eval()\n    all_outputs_C, all_outputs_F, all_outputs_P = [], [], []\n    \n    with torch.no_grad():\n        for features, _ in test_loader: # Chỉ unpack 2 giá trị\n            features = features.to(device)\n            output_C, output_F, output_P = model(features)\n            \n            all_outputs_C.append(output_C.cpu().numpy())\n            all_outputs_F.append(output_F.cpu().numpy())\n            all_outputs_P.append(output_P.cpu().numpy())\n            \n    pred_C = np.concatenate(all_outputs_C, axis=0)\n    pred_F = np.concatenate(all_outputs_F, axis=0)\n    pred_P = np.concatenate(all_outputs_P, axis=0)\n    \n    return pred_C, pred_F, pred_P\n\n\n# 2. CHUẨN BỊ TEST LOADER\ntest_X = np.stack(test_df['Emb'].values)\ntest_loader = DataLoader(CAFADatasetTest(test_X), batch_size=BATCH_SIZE, shuffle=False)\n\n# 3. LẤY LOGITS DỰ ĐOÁN\nprint(\"\\n--- DỰ ĐOÁN TRÊN TẬP TEST ---\")\ntest_pred_logits_C, test_pred_logits_F, test_pred_logits_P = get_test_predictions(\n    model, \n    test_loader, \n    DEVICE\n)\n\n\n# 4. TẠO SUBMISSION (LAN TRUYỀN VÀ CẮT NGƯỠNG)\nprint(\"--- BẮT ĐẦU LAN TRUYỀN (PROPAGATION) VÀ CẮT NGƯỠNG ---\")\nsubmission_records = []\ntest_entry_ids = test_df['EntryID'].values\n\n# Cần tạo đồ thị có hướng (GO Graph Directed) để truyền\n# Giả định bạn đã tạo biến này: go_graph_directed\n# Nếu chưa, bạn cần chạy: go_graph_directed = build_go_graph_directed(OBO_FILE_PATH, all_go_terms)\n\nfor aspect in aspects_list:\n    print(f\"  -> Xử lý Aspect {aspect}...\")\n    \n    # 1. Lấy dữ liệu\n    pred_logits = {'C': test_pred_logits_C, 'F': test_pred_logits_F, 'P': test_pred_logits_P}[aspect]\n    go_terms = aspect_data[aspect]['go_terms_all']\n    go_id_to_index = aspect_data[aspect]['go_id_to_index']\n    tau = 0.01\n    \n    # 2. Chuyển Logits sang Xác suất\n    pred_proba = 1.0 / (1.0 + np.exp(-pred_logits))\n    \n    # 3. LAN TRUYỀN XÁC SUẤT (PROPAGATE)\n    # LƯU Ý: Nếu go_graph_directed chưa được tạo, hàm propagate sẽ bị lỗi hoặc trả về kết quả sai.\n    propagated_proba = propagate_scores_to_parent(\n        pred_proba, go_graph_directed, go_id_to_index, go_terms\n    )\n    \n    # 4. Áp dụng Ngưỡng Cắt Tối ưu\n    test_predictions_masked = (propagated_proba >= tau)\n\n    # 5. Ghi kết quả vào Submission\n    for i, entry_id in enumerate(test_entry_ids):\n        predicted_indices = np.where(test_predictions_masked[i, :])[0]\n        \n        for idx in predicted_indices:\n            go_id = go_terms[idx]\n            score = propagated_proba[i, idx]\n            \n            # Chỉ nộp dự đoán có score > 0\n            if score > 0.0:\n                submission_records.append({\n                    'EntryID': entry_id,\n                    'GO_ID': go_id,\n                    'Score': score\n                })\n\n# 5. TẠO FILE SUBMISSION CUỐI CÙNG\nsubmission_df = pd.DataFrame(submission_records)\n\nsubmission_df = submission_df[['EntryID', 'GO_ID', 'Score']]\nsubmission_df = submission_df.sort_values(['EntryID', 'Score'], ascending=[True, False])\n\n\nprint(\"\\n--- HOÀN TẤT ---\")\nprint(\"Đã tạo file submission 'cafa6_final_submission.tsv'.\")\nprint(f\"Tổng số lượng dự đoán được tạo: {len(submission_df)}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:42:40.721652Z","iopub.execute_input":"2025-12-12T15:42:40.721883Z","iopub.status.idle":"2025-12-12T15:43:27.893128Z","shell.execute_reply.started":"2025-12-12T15:42:40.721861Z","shell.execute_reply":"2025-12-12T15:43:27.892471Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"submission_df.to_csv(\"submission.tsv\", sep=\"\\t\",index=False, header=False)","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-12T15:45:37.002963Z","iopub.execute_input":"2025-12-12T15:45:37.003389Z","iopub.status.idle":"2025-12-12T15:45:45.007915Z","shell.execute_reply.started":"2025-12-12T15:45:37.003356Z","shell.execute_reply":"2025-12-12T15:45:45.007322Z"}},"outputs":[],"execution_count":null}]}