{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"name":"python","version":"3.12.12","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,"isSourceIdPinned":false,"sourceType":"competition"},{"sourceId":14459097,"sourceType":"datasetVersion","datasetId":9235321}],"dockerImageVersionId":31234,"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 --upgrade gensim\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\nimport os\nimport gc\nfrom collections import defaultdict\n\nimport pandas as pd\nfrom tqdm.auto import tqdm\nimport numpy as np\n\n# ⚙️ CONFIGURATION - ADJUST THIS FOR SPEED VS ACCURACY\nSAMPLE_PERCENT = 100  # Use 100% of data\nQUICK_MODE = True   # Enable full feature computation\n\n# Package Installation\nimport subprocess\nimport sys\n\ndef install(package):\n    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"-q\", package])\n\nprint(\"Installing required packages...\")\ntry:\n    import obonet\nexcept:\n    install('obonet')\n    import obonet\n\ntry:\n    from Bio import SeqIO\nexcept:\n    install('biopython')\n    from Bio import SeqIO\n\n# Core Imports\nimport pandas as pd\nimport numpy as np\nfrom pathlib import Path\nfrom collections import defaultdict, Counter\nimport warnings\nwarnings.filterwarnings('ignore')\n\nimport networkx as nx\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.multiclass import OneVsRestClassifier\nfrom sklearn.linear_model import LogisticRegression\nimport Bio\n# Visualization imports\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom matplotlib.patches import Rectangle\nimport matplotlib.patches as mpatches\nplt.style.use('seaborn-v0_8-whitegrid')\nsns.set_palette(\"husl\")\n\nprint(\"=\"*80)\nprint(\"CAFA 6 PROTEIN FUNCTION PREDICTION - ENHANCED STARTER\")\nprint(f\"📊 SAMPLE MODE: {SAMPLE_PERCENT}% of data\")\nprint(f\"⚡ QUICK MODE: {'ON' if QUICK_MODE else 'OFF'}\")\nprint(\"=\"*80)\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"BASE_DIR = Path('/kaggle/input/cafa-6-protein-function-prediction')\nTRAIN_DIR = BASE_DIR / 'Train'\nTEST_DIR = BASE_DIR / 'Test'\ntrain_fasta = '/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta'\ntest_fasta = '/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta'\ntrain_obo = '/kaggle/input/cafa-6-protein-function-prediction/Train/go-basic.obo'\ntrain_text_path = '/kaggle/input/text-data/protein_text_features.csv'\ntest_text_path = '/kaggle/input/text-data/lrtest.csv'","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Train text embedding generation","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\n# Load your data from the .dat parser in the previous step\ndf_text = pd.read_csv(train_text_path)\n\n# Combine specifically: Name + Function + Literature Titles\n# We fill NaN with empty strings to avoid errors\ndf_text['description_combined'] = df_text['description'].fillna('') + \" \" + df_text['function'].fillna('')\ndf_text['literature_combined'] = df_text['literature'].fillna('')\n\n# The final 'narrative' used for embeddings\ndf_text['full_text'] = (df_text['description_combined'] + \" \" + df_text['literature_combined']).str.lower()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 2: TF-IDF (Term Importance)\nTF-IDF captures specific \"signal\" words (like kinase, zinc-finger, or autophagy) that appear in your protein’s literature but are rare across the other 82,000 proteins.","metadata":{}},{"cell_type":"code","source":"from sklearn.feature_extraction.text import TfidfVectorizer\n\n# NetGO logic uses a reduced dimensionality for TF-IDF to prevent overfitting\ntfidf_vectorizer = TfidfVectorizer(\n    max_features=1000,   # Keep the top 1000 most informative bio-terms\n    stop_words='english',\n    ngram_range=(1, 2)   # Capture both single words and pairs (e.g., 'binding site')\n)\n\nv_tfidf = tfidf_vectorizer.fit_transform(df_text['full_text']).toarray()\nprint(f\"TF-IDF Shape: {v_tfidf.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embed_dim = v_tfidf.shape[1]\n\ndf_TFIDF = pd.DataFrame(\n    v_tfidf,\n    columns=[f\"tfidf_{i}\" for i in range(embed_dim)]\n)\n\ndf_TFIDF.insert(0, \"protein_id\", df_text['protein_id'])\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_TFIDF.to_csv('tfidf_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 3: Doc2Vec (D2V - Global Semantic Theme)\nDoc2Vec represents the entire document as a fixed-length vector. While TF-IDF looks for specific words, D2V understands the overall \"theme\" of the description.","metadata":{}},{"cell_type":"code","source":"from gensim.models.doc2vec import Doc2Vec, TaggedDocument\n\n# 1. Tag each protein's text with its index\ndocuments = [TaggedDocument(doc.split(), [i]) for i, doc in enumerate(df_text['full_text'])]\n\n# 2. Train the model (256 is the standard dimension for D2V in AFP)\nd2v_model = Doc2Vec(documents, vector_size=256, window=8, min_count=2, workers=4, epochs=20)\n\n# 3. Infer vectors for all 82,000 proteins\nv_d2v = np.array([d2v_model.infer_vector(doc.split()) for doc in df_text['full_text']])\nprint(f\"Doc2Vec Shape: {v_d2v.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nembedding_dim = v_d2v.shape[1]\n\ndf_d2v = pd.DataFrame(\n    v_d2v,\n    columns=[f\"d2v_{i}\" for i in range(embedding_dim)]\n)\n\ndf_d2v.insert(0, \"protein_id\", df_text['protein_id'])\ndf_d2v","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_d2v.to_csv('d2v_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 4: PubMedBERT (Deep Biological Context)\nAs used in GORetrieval and GOXML, PubMedBERT provides a 768-dimensional vector that captures the deep scientific meaning of the literature.","metadata":{}},{"cell_type":"code","source":"import torch\nfrom transformers import AutoTokenizer, AutoModel\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel_name = \"microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext\"\n\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel = AutoModel.from_pretrained(model_name).to(device)\n\ndef get_pubmed_embeddings(texts, batch_size=64):\n    model.eval()\n    all_embs = []\n    for i in range(0, len(texts), batch_size):\n        batch = texts[i : i + batch_size]\n        inputs = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors=\"pt\").to(device)\n        \n        with torch.no_grad():\n            outputs = model(**inputs)\n            # Use the [CLS] token or Mean Pooling as the representation\n            embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n            all_embs.append(embeddings)\n    return np.vstack(all_embs)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"v_bert = get_pubmed_embeddings(df_text['full_text'].tolist())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"v_bert Shape: {v_bert.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embedding_pudmed = v_bert.shape[1]\n\ndf_pubmed = pd.DataFrame(\n    v_bert,\n    columns=[f\"medbert_{i}\" for i in range(embedding_pudmed)]\n)\n\ndf_pubmed.insert(0, \"protein_id\", df_text['protein_id'])\ndf_pubmed","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_pubmed.to_csv('pubmed_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Test embedding generation","metadata":{}},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\n\n# Load your data from the .dat parser in the previous step\ndf_test_text = pd.read_csv(test_text_path)\n\n# Combine specifically: Name + Function + Literature Titles\n# We fill NaN with empty strings to avoid errors\ndf_test_text['description_combined'] = df_test_text['description'].fillna('') + \" \" + df_test_text['function'].fillna('')\ndf_test_text['literature_combined'] = df_test_text['literature'].fillna('')\n\n# The final 'narrative' used for embeddings\ndf_test_text['full_text'] = (df_test_text['description_combined'] + \" \" + df_test_text['literature_combined']).str.lower()","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Step 2: TF-IDF (Term Importance)\nTF-IDF captures specific \"signal\" words (like kinase, zinc-finger, or autophagy) that appear in your protein’s literature but are rare across the other 82,000 proteins.¶¶","metadata":{}},{"cell_type":"code","source":"from sklearn.feature_extraction.text import TfidfVectorizer\n\n# NetGO logic uses a reduced dimensionality for TF-IDF to prevent overfitting\ntfidf_vectorizer = TfidfVectorizer(\n    max_features=1000,   # Keep the top 1000 most informative bio-terms\n    stop_words='english',\n    ngram_range=(1, 2)   # Capture both single words and pairs (e.g., 'binding site')\n)\n\ntest_tfidf = tfidf_vectorizer.fit_transform(df_test_text['full_text']).toarray()\nprint(f\"TF-IDF Shape: {test_tfidf.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embed_dim = test_tfidf.shape[1]\n\ndf_testTFIDF = pd.DataFrame(\n    test_tfidf,\n    columns=[f\"tfidf_{i}\" for i in range(embed_dim)]\n)\n\ndf_testTFIDF.insert(0, \"protein_id\", df_test_text['protein_id'])\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_testTFIDF.to_csv('test_tfidf_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"from gensim.models.doc2vec import Doc2Vec, TaggedDocument\n\n# 1. Tag each protein's text with its index\ndocuments = [TaggedDocument(doc.split(), [i]) for i, doc in enumerate(df_test_text['full_text'])]\n\n# 2. Train the model (256 is the standard dimension for D2V in AFP)\nd2v_model = Doc2Vec(documents, vector_size=256, window=8, min_count=2, workers=4, epochs=20)\n\n# 3. Infer vectors for all 82,000 proteins\ntest_d2v = np.array([d2v_model.infer_vector(doc.split()) for doc in df_test_text['full_text']])\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"print(f\"Doc2Vec Shape: {test_d2v.shape}\")","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\n\nembedding_dim = test_d2v.shape[1]\n\ndf_testd2v = pd.DataFrame(\n    test_d2v,\n    columns=[f\"d2v_{i}\" for i in range(embedding_dim)]\n)\n\ndf_testd2v.insert(0, \"protein_id\", df_test_text['protein_id'])\ndf_testd2v","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_testd2v.to_csv('d2vtest_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom transformers import AutoTokenizer, AutoModel\n\ndevice = \"cuda\" if torch.cuda.is_available() else \"cpu\"\nmodel_name = \"microsoft/BiomedNLP-PubMedBERT-base-uncased-abstract-fulltext\"\n\ntokenizer = AutoTokenizer.from_pretrained(model_name)\nmodel = AutoModel.from_pretrained(model_name).to(device)\n\ndef get_pubmed_embeddings(texts, batch_size=64):\n    model.eval()\n    all_embs = []\n    for i in range(0, len(texts), batch_size):\n        batch = texts[i : i + batch_size]\n        inputs = tokenizer(batch, padding=True, truncation=True, max_length=512, return_tensors=\"pt\").to(device)\n        \n        with torch.no_grad():\n            outputs = model(**inputs)\n            # Use the [CLS] token or Mean Pooling as the representation\n            embeddings = outputs.last_hidden_state.mean(dim=1).cpu().numpy()\n            all_embs.append(embeddings)\n    return np.vstack(all_embs)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"v_bert = get_pubmed_embeddings(df_test_text['full_text'].tolist())","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"embedding_pudmed = v_bert.shape[1]\n\ndf_testpubmed = pd.DataFrame(\n    v_bert,\n    columns=[f\"medbert_{i}\" for i in range(embedding_pudmed)]\n)\n\ndf_testpubmed.insert(0, \"protein_id\", df_test_text['protein_id'])\ndf_testpubmed","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"code","source":"df_testpubmed.to_csv('test_pubmed_embedding.csv', index=False)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}