{"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":"none","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"},{"sourceId":616258,"sourceType":"modelInstanceVersion","isSourceIdPinned":true,"modelInstanceId":463229,"modelId":479009}],"dockerImageVersionId":31192,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q faiss-cpu\n!pip install biopython\n!pip install protobuf==4.25.3 --force-reinstall\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:42:45.909626Z","iopub.execute_input":"2025-12-02T04:42:45.909963Z","iopub.status.idle":"2025-12-02T04:42:59.147607Z","shell.execute_reply.started":"2025-12-02T04:42:45.909938Z","shell.execute_reply":"2025-12-02T04:42:59.145961Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom pathlib import Path\nfrom Bio import SeqIO\nimport matplotlib.pyplot as plt\nimport os, sys\nprint(\"Reinstall done. Please restart the kernel (or this process will exit).\")\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.decomposition import PCA\nfrom sklearn.cluster import KMeans\nfrom collections import Counter\n\nDATA_ROOT = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\n\n# ---- train_terms.tsv ----\n# 보통 컬럼: EntryID, term, [aspect(optional)]\ntrain_terms = pd.read_csv(DATA_ROOT / \"Train/train_terms.tsv\", sep=\"\\t\")\n\nprint(train_terms.head())\nprint(train_terms.columns)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:42:59.149579Z","iopub.execute_input":"2025-12-02T04:42:59.149883Z","iopub.status.idle":"2025-12-02T04:42:59.465139Z","shell.execute_reply.started":"2025-12-02T04:42:59.149855Z","shell.execute_reply":"2025-12-02T04:42:59.463948Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- train_sequences.fasta ----\nids = []\nseqs = []\nwith open(DATA_ROOT / \"Train/train_sequences.fasta\") as handle:\n    for record in SeqIO.parse(handle, \"fasta\"):\n        ids.append(record.id)\n        seqs.append(str(record.seq))\n\ntrain_seqs_df = pd.DataFrame({\"EntryID\": ids, \"Seq\": seqs})\ntrain_seqs_df[\"seq_len\"] = train_seqs_df[\"Seq\"].str.len()\n\nprint(train_seqs_df.head())\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:42:59.466092Z","iopub.execute_input":"2025-12-02T04:42:59.466411Z","iopub.status.idle":"2025-12-02T04:43:00.068367Z","shell.execute_reply.started":"2025-12-02T04:42:59.466362Z","shell.execute_reply":"2025-12-02T04:43:00.067345Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# ---- terms를 protein 기준으로 집계 ----\n# 하나의 EntryID에 여러 term이 붙어 있으므로 groupby\nterms_grouped = (\n    train_terms\n    .groupby(\"EntryID\")[\"term\"]\n    .apply(list)\n    .reset_index()\n    .rename(columns={\"term\": \"GO_terms\"})\n)\n\n# aspect 컬럼이 있으면 같이 모아두기 (MF/BP/CC)\nif \"aspect\" in train_terms.columns:\n    aspect_grouped = (\n        train_terms\n        .groupby(\"EntryID\")[\"aspect\"]\n        .apply(lambda x: list(set(x)))\n        .reset_index()\n        .rename(columns={\"aspect\": \"aspects\"})\n    )\n    train_info = (\n        train_seqs_df\n        .merge(terms_grouped, on=\"EntryID\", how=\"left\")\n        .merge(aspect_grouped, on=\"EntryID\", how=\"left\")\n    )\nelse:\n    train_info = (\n        train_seqs_df\n        .merge(terms_grouped, on=\"EntryID\", how=\"left\")\n    )\n\ntrain_info[\"n_terms\"] = train_info[\"GO_terms\"].apply(lambda x: len(x) if isinstance(x, list) else 0)\nprint(train_info.head())\nprint(train_info.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:00.071348Z","iopub.execute_input":"2025-12-02T04:43:00.071832Z","iopub.status.idle":"2025-12-02T04:43:04.000876Z","shell.execute_reply.started":"2025-12-02T04:43:00.071802Z","shell.execute_reply":"2025-12-02T04:43:03.999362Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 1) 시퀀스 길이 히스토그램\nplt.figure(figsize=(6,4))\nplt.hist(train_info[\"seq_len\"], bins=50)\nplt.xlabel(\"Sequence length (aa)\")\nplt.ylabel(\"Count\")\nplt.title(\"Distribution of protein sequence lengths (Train)\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:04.001848Z","iopub.execute_input":"2025-12-02T04:43:04.002184Z","iopub.status.idle":"2025-12-02T04:43:04.263414Z","shell.execute_reply.started":"2025-12-02T04:43:04.002157Z","shell.execute_reply":"2025-12-02T04:43:04.262316Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2) 한 단백질당 GO term 개수 분포\nplt.figure(figsize=(6,4))\n# 너무 헤비하면 상한 잘라서 보기\nmax_terms_to_show = 50\nvals = np.clip(train_info[\"n_terms\"], 0, max_terms_to_show)\nplt.hist(vals, bins=min(max_terms_to_show, 50))\nplt.xlabel(f\"#GO terms per protein (clipped at {max_terms_to_show})\")\nplt.ylabel(\"Count\")\nplt.title(\"Distribution of #GO terms per protein\")\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:04.264449Z","iopub.execute_input":"2025-12-02T04:43:04.264716Z","iopub.status.idle":"2025-12-02T04:43:04.510653Z","shell.execute_reply.started":"2025-12-02T04:43:04.264696Z","shell.execute_reply":"2025-12-02T04:43:04.509628Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 3) aspect별 단백질 수 (있을 때만)\nif \"aspects\" in train_info.columns:\n    all_aspects = []\n    for lst in train_info[\"aspects\"].dropna():\n        all_aspects.extend(lst)\n    cnt = Counter(all_aspects)\n    print(\"Aspect counts (MF/BP/CC):\", cnt)\n\n    plt.figure(figsize=(4,4))\n    plt.bar(list(cnt.keys()), list(cnt.values()))\n    plt.xlabel(\"GO sub-ontology (aspect)\")\n    plt.ylabel(\"#Proteins (train)\")\n    plt.title(\"Number of proteins per GO aspect\")\n    plt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:04.511808Z","iopub.execute_input":"2025-12-02T04:43:04.512166Z","iopub.status.idle":"2025-12-02T04:43:04.680409Z","shell.execute_reply.started":"2025-12-02T04:43:04.512136Z","shell.execute_reply":"2025-12-02T04:43:04.679477Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 4) 가장 많이 등장한 GO term Top 20\nall_terms = []\nfor lst in train_info[\"GO_terms\"].dropna():\n    all_terms.extend(lst)\nterm_cnt = Counter(all_terms)\ntop_terms = term_cnt.most_common(20)\ntop_terms_df = pd.DataFrame(top_terms, columns=[\"GO_term\", \"count\"])\ndisplay(top_terms_df)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:04.681343Z","iopub.execute_input":"2025-12-02T04:43:04.681620Z","iopub.status.idle":"2025-12-02T04:43:04.697598Z","shell.execute_reply.started":"2025-12-02T04:43:04.681600Z","shell.execute_reply":"2025-12-02T04:43:04.696446Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 아미노산 20종 (표준)\nAA_LIST = list(\"ACDEFGHIKLMNPQRSTVWY\")\n\ndef aa_composition(seq):\n    total = len(seq)\n    counts = np.zeros(len(AA_LIST), dtype=float)\n    if total == 0:\n        return counts\n    for i, aa in enumerate(AA_LIST):\n        counts[i] = seq.count(aa)\n    return counts / total\n\n# feature matrix 만들기\nX_aa = np.vstack(train_info[\"Seq\"].apply(aa_composition).values)\nprint(\"Feature matrix shape:\", X_aa.shape)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:04.698798Z","iopub.execute_input":"2025-12-02T04:43:04.699109Z","iopub.status.idle":"2025-12-02T04:43:06.640537Z","shell.execute_reply.started":"2025-12-02T04:43:04.699078Z","shell.execute_reply":"2025-12-02T04:43:06.639235Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 스케일링 + PCA 2D\nscaler = StandardScaler()\nX_scaled = scaler.fit_transform(X_aa)\n\npca = PCA(n_components=2, random_state=42)\nX_pca = pca.fit_transform(X_scaled)\nprint(\"Explained variance ratio:\", pca.explained_variance_ratio_)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:06.643665Z","iopub.execute_input":"2025-12-02T04:43:06.643958Z","iopub.status.idle":"2025-12-02T04:43:06.872698Z","shell.execute_reply.started":"2025-12-02T04:43:06.643938Z","shell.execute_reply":"2025-12-02T04:43:06.871196Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# KMeans 클러스터 (예: 10개 클러스터, 필요하면 조정)\nn_clusters = 10\nkmeans = KMeans(n_clusters=n_clusters, random_state=42, n_init=\"auto\")\ncluster_labels = kmeans.fit_predict(X_scaled)\n\ntrain_info[\"cluster\"] = cluster_labels\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:06.873394Z","iopub.execute_input":"2025-12-02T04:43:06.873683Z","iopub.status.idle":"2025-12-02T04:43:07.462947Z","shell.execute_reply.started":"2025-12-02T04:43:06.873659Z","shell.execute_reply":"2025-12-02T04:43:07.461986Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 2D PCA + cluster 색\nplt.figure(figsize=(6,6))\nfor c in range(n_clusters):\n    mask = (train_info[\"cluster\"] == c)\n    plt.scatter(\n        X_pca[mask, 0],\n        X_pca[mask, 1],\n        s=5,\n        alpha=0.5,\n        label=f\"cluster {c}\",\n    )\nplt.xlabel(\"PCA 1\")\nplt.ylabel(\"PCA 2\")\nplt.title(\"Amino-acid composition clusters (KMeans on Train)\")\nplt.legend(markerscale=3, fontsize=8, bbox_to_anchor=(1.05, 1), loc=\"upper left\")\nplt.tight_layout()\nplt.show()\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:07.463874Z","iopub.execute_input":"2025-12-02T04:43:07.464164Z","iopub.status.idle":"2025-12-02T04:43:08.018433Z","shell.execute_reply.started":"2025-12-02T04:43:07.464138Z","shell.execute_reply":"2025-12-02T04:43:08.017572Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 클러스터별 기본 통계\ncluster_stats = (\n    train_info\n    .groupby(\"cluster\")\n    .agg(\n        n_proteins=(\"EntryID\", \"nunique\"),\n        mean_seq_len=(\"seq_len\", \"mean\"),\n        median_seq_len=(\"seq_len\", \"median\"),\n        mean_n_terms=(\"n_terms\", \"mean\"),\n    )\n    .reset_index()\n    .sort_values(\"cluster\")\n)\n\ndisplay(cluster_stats)\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:08.019354Z","iopub.execute_input":"2025-12-02T04:43:08.019661Z","iopub.status.idle":"2025-12-02T04:43:08.067748Z","shell.execute_reply.started":"2025-12-02T04:43:08.019633Z","shell.execute_reply":"2025-12-02T04:43:08.066913Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"# 각 클러스터별 상위 GO term 5개 뽑기\ncluster_top_terms = {}\n\nfor c in range(n_clusters):\n    sub = train_info[train_info[\"cluster\"] == c]\n    all_terms = []\n    for lst in sub[\"GO_terms\"].dropna():\n        all_terms.extend(lst)\n    cnt = Counter(all_terms)\n    cluster_top_terms[c] = cnt.most_common(5)\n\nfor c in range(n_clusters):\n    print(f\"\\n=== Cluster {c} ===\")\n    print(\"n_proteins:\", cluster_stats.loc[cluster_stats[\"cluster\"]==c, \"n_proteins\"].values[0])\n    print(\"Top GO terms (term, count):\")\n    for term, count in cluster_top_terms[c]:\n        print(f\"  {term}: {count}\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:08.068827Z","iopub.execute_input":"2025-12-02T04:43:08.069096Z","iopub.status.idle":"2025-12-02T04:43:08.116876Z","shell.execute_reply.started":"2025-12-02T04:43:08.069075Z","shell.execute_reply":"2025-12-02T04:43:08.115600Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom pathlib import Path\nimport matplotlib.pyplot as plt\nfrom collections import Counter\n\nDATA_ROOT = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\n\n# ---- 1) Biopython 없이 FASTA 읽기 ----\ndef read_fasta_as_df(path: Path):\n    ids, seqs = [], []\n    cur_id, cur_seq = None, []\n    with open(path, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                if cur_id is not None:\n                    ids.append(cur_id)\n                    seqs.append(\"\".join(cur_seq))\n                cur_id = line[1:].split()[0]\n                cur_seq = []\n            else:\n                cur_seq.append(line)\n    if cur_id is not None:\n        ids.append(cur_id)\n        seqs.append(\"\".join(cur_seq))\n    return pd.DataFrame({\"EntryID\": ids, \"Seq\": seqs})\n\n# ---- 2) train 시퀀스 / 라벨 로드 ----\ntrain_seq_df = read_fasta_as_df(DATA_ROOT / \"Train/train_sequences.fasta\")\ntrain_seq_df[\"seq_len\"] = train_seq_df[\"Seq\"].str.len()\n\ntrain_terms = pd.read_csv(DATA_ROOT / \"Train/train_terms.tsv\", sep=\"\\t\")\n# 한 단백질당 라벨 개수\nterms_per_protein = train_terms.groupby(\"EntryID\")[\"term\"].nunique()\n\n# ---- 3) 아주 기본적인 통계 ----\nprint(\"Train 단백질 수:\", len(train_seq_df))\nprint(\"GO term 라벨 수(유니크):\", train_terms[\"term\"].nunique())\nprint(\"단백질별 라벨 개수 평균:\", terms_per_protein.mean())\nprint(\"단백질별 라벨 개수 중간값:\", terms_per_protein.median())\n\n# ---- 4) 시각화 1: 시퀀스 길이 ----\nplt.figure(figsize=(6,4))\nplt.hist(train_seq_df[\"seq_len\"], bins=50)\nplt.xlabel(\"Sequence length\")\nplt.ylabel(\"Count\")\nplt.title(\"Train sequence length distribution\")\nplt.show()\n\n# ---- 5) 시각화 2: 라벨 개수 ----\nplt.figure(figsize=(6,4))\nplt.hist(terms_per_protein.clip(upper=20), bins=20)\nplt.xlabel(\"#GO terms per protein (capped at 20)\")\nplt.ylabel(\"Count\")\nplt.title(\"Labels per protein\")\nplt.show()\n\n# ---- 6) 통계 3: 자주 나오는 GO term TOP 30 ----\nterm_counts = train_terms[\"term\"].value_counts()\nprint(\"\\nTop 30 frequent GO terms:\")\nprint(term_counts.head(30))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:08.117670Z","iopub.execute_input":"2025-12-02T04:43:08.117910Z","iopub.status.idle":"2025-12-02T04:43:09.573170Z","shell.execute_reply.started":"2025-12-02T04:43:08.117892Z","shell.execute_reply":"2025-12-02T04:43:09.572096Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\n\n# =========================\n# 기본 설정\n# =========================\nDATA_ROOT = Path(\"/kaggle/input/cafa-6-protein-function-prediction\")\nMODEL_DIR = \"/kaggle/input/esm2-150m/transformers/default/1/esm2_150M_local\"\n\nprint(\"DATA_ROOT:\", DATA_ROOT)\nprint(\"MODEL_DIR:\", MODEL_DIR)\n\n# =========================\n# FASTA → DataFrame 변환 함수\n# =========================\ndef read_fasta_to_df(path_fasta: Path, id_col: str = \"EntryID\", seq_col: str = \"Seq\"):\n    ids = []\n    seqs = []\n    cur_id = None\n    cur_seq = []\n    with open(path_fasta, \"r\") as f:\n        for line in f:\n            line = line.strip()\n            if not line:\n                continue\n            if line.startswith(\">\"):\n                # 새로운 시퀀스 시작 → 이전 것 저장\n                if cur_id is not None:\n                    ids.append(cur_id)\n                    seqs.append(\"\".join(cur_seq))\n                header = line[1:]\n                # 헤더의 첫 토큰을 ID로 사용 (UniProt ID)\n                cur_id = header.split()[0]\n                cur_seq = []\n            else:\n                cur_seq.append(line)\n        # 마지막 시퀀스\n        if cur_id is not None:\n            ids.append(cur_id)\n            seqs.append(\"\".join(cur_seq))\n    df = pd.DataFrame({id_col: ids, seq_col: seqs})\n    return df\n\n# =========================\n# 1) Train / Test 데이터 로드\n# =========================\n# train_terms: (EntryID, term)\ntrain_terms = pd.read_csv(DATA_ROOT / \"Train/train_terms.tsv\", sep=\"\\t\")\nprint(\"train_terms.shape:\", train_terms.shape)\n\n# train 시퀀스\ntrain_seq_df = read_fasta_to_df(DATA_ROOT / \"Train/train_sequences.fasta\")\nprint(\"train_seq_df.shape:\", train_seq_df.shape)\n\n# test superset 시퀀스 (경로: /Test/testsuperset.fasta)\ntest_seq_df = read_fasta_to_df(DATA_ROOT / \"Test/testsuperset.fasta\")\nprint(\"test_seq_df.shape:\", test_seq_df.shape)\n\n# =========================\n# 2) Protein ID / GO term 인덱싱\n# =========================\n# 훈련에 사용할 protein 리스트\ntrain_ids = sorted(train_seq_df[\"EntryID\"].unique())\ntest_ids  = sorted(test_seq_df[\"EntryID\"].unique())\n\nid2idx_train = {pid: i for i, pid in enumerate(train_ids)}\nid2idx_test  = {pid: i for i, pid in enumerate(test_ids)}\n\n# GO term 전체 set\nall_terms = sorted(train_terms[\"term\"].unique())\nterm2idx = {t: i for i, t in enumerate(all_terms)}\nidx2term = {i: t for t, i in term2idx.items()}\n\nn_train = len(train_ids)\nn_test  = len(test_ids)\nn_labels = len(all_terms)\n\nprint(f\"#train proteins: {n_train}, #test proteins: {n_test}, #GO terms: {n_labels}\")\n\n# =========================\n# 3) 멀티라벨 타겟 매트릭스 Y (multi-hot)\n# =========================\nY = np.zeros((n_train, n_labels), dtype=np.float32)\n\nfor entry_id, term in tqdm(\n    train_terms[[\"EntryID\", \"term\"]].itertuples(index=False),\n    total=len(train_terms),\n    desc=\"Building label matrix\",\n):\n    if entry_id not in id2idx_train:\n        # 혹시 train_sequences에 없는 ID가 있으면 skip\n        continue\n    i = id2idx_train[entry_id]\n    j = term2idx[term]\n    Y[i, j] = 1.0\n\nprint(\"Y.shape:\", Y.shape)\n\n# =========================\n# 4) 시퀀스 리스트 (train / test) — 빠른 버전\n# =========================\nprint(\"[STEP 4] Building sequence lists with fast mapping...\")\n\n# 1) EntryID → Seq 딕셔너리로 한 번에 변환\ntrain_seq_map = dict(zip(train_seq_df[\"EntryID\"].values,\n                         train_seq_df[\"Seq\"].values))\ntest_seq_map  = dict(zip(test_seq_df[\"EntryID\"].values,\n                         test_seq_df[\"Seq\"].values))\n\nprint(\"       train_seq_map size:\", len(train_seq_map))\nprint(\"       test_seq_map  size:\", len(test_seq_map))\n\n# 2) train_ids / test_ids 순서에 맞춰 리스트 생성\n#    (이건 O(N) 이라 금방 끝남)\nmissing_train = [pid for pid in train_ids if pid not in train_seq_map]\nmissing_test  = [pid for pid in test_ids  if pid not in test_seq_map]\n\nif missing_train:\n    print(\"[WARN] Some train_ids are missing in train_seq_df! e.g.:\", missing_train[:5])\nif missing_test:\n    print(\"[WARN] Some test_ids are missing in test_seq_df! e.g.:\", missing_test[:5])\n\ntrain_seqs = [train_seq_map[pid] for pid in train_ids]\ntest_seqs  = [test_seq_map[pid]  for pid in test_ids]\n\nprint(\"len(train_seqs):\", len(train_seqs))\nprint(\"len(test_seqs):\", len(test_seqs))\n\n# sanity check\nprint(\"예시 train 시퀀스:\",\n      train_ids[0],\n      train_seqs[0][:50] + (\"...\" if len(train_seqs[0]) > 50 else \"\"))\nprint(\"예시 test  시퀀스:\",\n      test_ids[0],\n      test_seqs[0][:50] + (\"...\" if len(test_seqs[0]) > 50 else \"\"))\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:09.574157Z","iopub.execute_input":"2025-12-02T04:43:09.574990Z","iopub.status.idle":"2025-12-02T04:43:12.288440Z","shell.execute_reply.started":"2025-12-02T04:43:09.574952Z","shell.execute_reply":"2025-12-02T04:43:12.287225Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\n\nDEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nprint(\"DEVICE:\", DEVICE)\nprint(\"MODEL_DIR:\", MODEL_DIR)\n\n# 앞 셀에서 이미 만들어진 것들:\n# - train_seqs, test_seqs\n# - train_ids, test_ids\nfor v in [\"train_seqs\", \"test_seqs\", \"train_ids\", \"test_ids\"]:\n    assert v in globals(), f\"{v} 가 정의돼 있어야 합니다. 앞 셀을 먼저 실행하세요.\"\n\n# ------------------------------------------------\n# 1) ESM2 로딩\n# ------------------------------------------------\nprint(\"[ESM] Loading tokenizer/model...\")\ntokenizer = AutoTokenizer.from_pretrained(MODEL_DIR)\nesm_model = AutoModel.from_pretrained(MODEL_DIR)\nesm_model.to(DEVICE)\nesm_model.eval()\n\n# ------------------------------------------------\n# 2) 시퀀스 → 임베딩 함수\n# ------------------------------------------------\ndef embed_sequences_with_esm2(seqs, batch_size=8, max_length=1022, desc=\"ESM2\"):\n    \"\"\"\n    seqs: list[str]\n    return: np.ndarray, shape = (len(seqs), hidden_dim)\n    \"\"\"\n    all_embeds = []\n    for i in tqdm(range(0, len(seqs), batch_size), desc=desc):\n        batch_seqs = seqs[i : i + batch_size]\n\n        with torch.no_grad():\n            enc = tokenizer(\n                batch_seqs,\n                return_tensors=\"pt\",\n                padding=True,\n                truncation=True,\n                max_length=max_length,\n            )\n            enc = {k: v.to(DEVICE) for k, v in enc.items()}\n            out = esm_model(**enc)\n            hidden = out.last_hidden_state  # (B, L, D)\n\n            # attention_mask를 이용해서 mean pooling\n            attn = enc[\"attention_mask\"].unsqueeze(-1)  # (B, L, 1)\n            masked_hidden = hidden * attn\n            sum_hidden = masked_hidden.sum(dim=1)       # (B, D)\n            lengths = attn.sum(dim=1).clamp(min=1)      # (B, 1)\n            emb = (sum_hidden / lengths).cpu().numpy()  # (B, D)\n\n        all_embeds.append(emb)\n\n    embs = np.concatenate(all_embeds, axis=0).astype(np.float32)\n    return embs\n\n# ------------------------------------------------\n# 3) 캐싱 경로 설정\n# ------------------------------------------------\nwork_dir = Path(\"/kaggle/working\")\nwork_dir.mkdir(parents=True, exist_ok=True)\n\ntrain_emb_path = work_dir / \"train_embeddings_esm2_150m.npy\"\ntest_emb_path  = work_dir / \"test_embeddings_esm2_150m.npy\"\n\n# ------------------------------------------------\n# 4) 이미 있으면 로드, 없으면 계산 후 저장\n# ------------------------------------------------\nif train_emb_path.exists() and test_emb_path.exists():\n    print(\"[ESM] Found cached embeddings. Skipping computation.\")\n    train_embs = np.load(train_emb_path)\n    test_embs  = np.load(test_emb_path)\n    print(\"       train_embs.shape:\", train_embs.shape)\n    print(\"       test_embs.shape :\", test_embs.shape)\nelse:\n    print(\"[ESM] No cached embeddings found. Computing TRAIN embeddings...\")\n    train_embs = embed_sequences_with_esm2(\n        train_seqs,\n        batch_size=8,      # GPU 여유에 따라 4~16 사이 조절\n        max_length=1022,\n        desc=\"ESM2 train\"\n    )\n    print(\"       train_embs.shape:\", train_embs.shape)\n    np.save(train_emb_path, train_embs)\n    print(\"       Saved:\", train_emb_path)\n\n    print(\"[ESM] Computing TEST embeddings...\")\n    test_embs = embed_sequences_with_esm2(\n        test_seqs,\n        batch_size=8,\n        max_length=1022,\n        desc=\"ESM2 test\"\n    )\n    print(\"       test_embs.shape:\", test_embs.shape)\n    np.save(test_emb_path, test_embs)\n    print(\"       Saved:\", test_emb_path)\n\n# 메모리 정리 (원하면 여기서 바로 FAISS에 넘겨도 됨)\ndel train_embs, test_embs\ngc.collect()\n\nprint(\"[ESM] Embedding step finished.\")\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-02T04:43:12.289659Z","iopub.execute_input":"2025-12-02T04:43:12.290448Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import gc\nimport numpy as np\nimport pandas as pd\nfrom pathlib import Path\nfrom tqdm.auto import tqdm\n\n# FAISS (CPU 버전)\nimport faiss\n\n# ============================================\n# 0. 전제: 앞 셀에서 이미 정의된 것들\n#    - DATA_ROOT, MODEL_DIR\n#    - train_ids, test_ids\n#    - Y (shape: [n_train, n_labels])\n#    - idx2term\n# ============================================\nfor v in [\"train_ids\", \"test_ids\", \"Y\", \"idx2term\"]:\n    assert v in globals(), f\"{v} 가 정의돼 있어야 합니다. 앞 셀을 먼저 실행하세요.\"\n\nn_train = len(train_ids)\nn_test  = len(test_ids)\nn_labels = Y.shape[1]\n\nprint(f\"[INFO] #train={n_train}, #test={n_test}, #labels={n_labels}\")\n\n# ============================================\n# 1. ESM 임베딩 로드 (이미 만들어진 임베딩을 그대로 사용)\n# ============================================\nwork_dir = Path(\"/kaggle/working\")\ntrain_emb_path = work_dir / \"train_embeddings_esm2_150m.npy\"\ntest_emb_path  = work_dir / \"test_embeddings_esm2_150m.npy\"\n\nassert train_emb_path.exists(), f\"{train_emb_path} 가 없습니다. ESM 임베딩 생성 셀을 먼저 실행하세요.\"\nassert test_emb_path.exists(),  f\"{test_emb_path} 가 없습니다. ESM 임베딩 생성 셀을 먼저 실행하세요.\"\n\nprint(\"[STEP 1] Loading ESM2 embeddings from disk...\")\nX_train = np.load(train_emb_path).astype(np.float32)\nX_test  = np.load(test_emb_path).astype(np.float32)\n\nprint(\"       X_train.shape:\", X_train.shape)\nprint(\"       X_test.shape :\", X_test.shape)\n\n# ============================================\n# 2. L2 정규화 (cosine 유사도 ≒ inner product)\n# ============================================\ndef l2_normalize(x: np.ndarray) -> np.ndarray:\n    norms = np.linalg.norm(x, axis=1, keepdims=True)\n    norms[norms == 0] = 1.0\n    return x / norms\n\nprint(\"[STEP 2] L2-normalizing embeddings...\")\nX_train_norm = l2_normalize(X_train)\nX_test_norm  = l2_normalize(X_test)\n\ndel X_train, X_test\ngc.collect()\n\n# ============================================\n# 3. FAISS 인덱스 구성 (IndexFlatIP: inner product)\n# ============================================\nd = X_train_norm.shape[1]\nprint(f\"[STEP 3] Building FAISS index (dim={d})...\")\n\nindex = faiss.IndexFlatIP(d)\nindex.add(X_train_norm)\nprint(\"       Indexed vectors:\", index.ntotal)\n\nK = 64  # 32~128 사이에서 조절 가능\nprint(\"       kNN neighbors K =\", K)\n\n# ============================================\n# 4. FAISS 검색을 chunk 로 나눠서 진행상황 보기\n# ============================================\nprint(\"[STEP 4] Running FAISS search over test set (chunked)...\")\n\nBATCH_SEARCH = 1024  # test를 나눠서 탐색 (너무 크면 메모리↑, 너무 작으면 느려짐)\nD_list = []\nI_list = []\n\nfor start in tqdm(range(0, n_test, BATCH_SEARCH), desc=\"FAISS search (test)\"):\n    end = min(n_test, start + BATCH_SEARCH)\n    # X_test_norm[start:end]에 대해 kNN 검색\n    Di, Ii = index.search(X_test_norm[start:end], K)\n    D_list.append(Di)\n    I_list.append(Ii)\n\nD = np.vstack(D_list)\nI = np.vstack(I_list)\n\nprint(\"       D.shape:\", D.shape, \"I.shape:\", I.shape)\n\n# ============================================\n# 5. train protein별 label 인덱스 리스트 (희소 표현)\n# ============================================\nprint(\"[STEP 5] Building sparse label lists for train proteins...\")\ntrain_label_lists = []\nfor i in tqdm(range(n_train), desc=\"train_label_lists\"):\n    lbl_idx = np.where(Y[i] > 0)[0]\n    train_label_lists.append(lbl_idx)\n\n# ============================================\n# 6. kNN 기반 label propagation\n# ============================================\nMIN_PROB = 0.01           # 이 값 이상인 label만 예측으로 사용\nMAX_TERMS_PER_PROT = 300  # protein당 최대 term 수 (파일 사이즈 제어)\n\nrows = []\n\nprint(\"[STEP 6] Building kNN-based submission rows...\")\nfor test_idx, pid in tqdm(\n    enumerate(test_ids),\n    total=n_test,\n    desc=\"test proteins\"\n):\n    neighs = I[test_idx]        # shape (K,)\n    sims  = D[test_idx].copy()  # shape (K,)\n\n    # 음수 유사도 제거 (0으로 취급) → \"비슷하지 않다\"\n    sims = np.maximum(sims, 0.0)\n\n    if sims.sum() == 0:\n        sims[:] = 1.0 / len(sims)\n    else:\n        sims /= sims.sum()\n\n    label_scores = {}\n\n    # 각 이웃 protein이 가진 label에 weight를 누적\n    for w, train_idx in zip(sims, neighs):\n        if w <= 0:\n            continue\n        for lbl in train_label_lists[train_idx]:\n            label_scores[lbl] = label_scores.get(lbl, 0.0) + float(w)\n\n    if not label_scores:\n        continue\n\n    # score 내림차순 정렬 후 상위 label들만 사용\n    sorted_labels = sorted(label_scores.items(), key=lambda x: -x[1])\n\n    count = 0\n    for lbl_idx, score in sorted_labels:\n        if score < MIN_PROB:\n            break\n        term = idx2term[lbl_idx]\n        score_clamped = max(min(score, 1.0), 1e-5)\n        rows.append((pid, term, score_clamped))\n        count += 1\n        if count >= MAX_TERMS_PER_PROT:\n            break\n\n# ============================================\n# 7. DataFrame → submission_knn.tsv 저장\n# ============================================\nprint(\"[STEP 7] Saving submission_knn.tsv ...\")\nsubmission_knn = pd.DataFrame(rows, columns=[\"Id\", \"GO term\", \"Confidence\"])\nsubmission_knn.to_csv(\"submission_knn.tsv\", sep=\"\\t\", header=False, index=False)\n\nprint(\"       Saved submission_knn.tsv with\", len(submission_knn), \"rows\")\nsubmission_knn.head()\n","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}