{"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":"gpu","dataSources":[{"sourceId":116062,"databundleVersionId":14084779,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install -q torch transformers sentencepiece biopython pandas tqdm ","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T16:17:04.466558Z","iopub.execute_input":"2025-12-09T16:17:04.466817Z","iopub.status.idle":"2025-12-09T16:18:12.237185Z","shell.execute_reply.started":"2025-12-09T16:17:04.466796Z","shell.execute_reply":"2025-12-09T16:18:12.236438Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import os\n\nos.environ[\"USE_TORCH\"] = \"1\"\nos.environ[\"USE_TF\"] = \"0\"\nos.environ[\"TRANSFORMERS_NO_ADVISORY_WARNINGS\"] = \"true\"","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T16:18:12.238690Z","iopub.execute_input":"2025-12-09T16:18:12.238889Z","iopub.status.idle":"2025-12-09T16:18:12.243008Z","shell.execute_reply.started":"2025-12-09T16:18:12.238869Z","shell.execute_reply":"2025-12-09T16:18:12.242267Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import torch\nfrom transformers import AutoTokenizer, AutoModel\nfrom Bio import SeqIO\nimport numpy as np\nimport pandas as pd\nimport os\nimport gc\nfrom tqdm.auto import tqdm\nimport sys\n\n# ==========================================\n# 1. CẤU HÌNH (CONFIGURATION)\n# ==========================================\nclass CFG:\n    # Đường dẫn input (Raw data)\n    DATA_DIR = \"/kaggle/input/cafa-6-protein-function-prediction\" \n    TRAIN_FASTA = os.path.join(DATA_DIR, \"Train/train_sequences.fasta\")\n    TEST_FASTA = os.path.join(DATA_DIR, \"Test/testsuperset.fasta\")\n    \n    # Đường dẫn output (Nơi lưu file .npy)\n    OUTPUT_DIR = \"/kaggle/working\"\n    \n    MODEL_NAME = \"facebook/esm2_t33_650M_UR50D\"\n    \n    # Tham số chạy\n    MAX_LEN = 1024      # Cắt các chuỗi dài hơn 1024 aa\n    BATCH_SIZE = 32     # Tăng lên 64 nếu GPU khỏe, giảm xuống 16 nếu OOM\n    DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n\nprint(f\"Device: {CFG.DEVICE}\")\nprint(f\"Model: {CFG.MODEL_NAME}\")\n\n# ==========================================\n# 2. HÀM XỬ LÝ (CORE FUNCTIONS)\n# ==========================================\n\ndef load_model():\n    print(\"Downloading & Loading Model...\")\n    tokenizer = AutoTokenizer.from_pretrained(CFG.MODEL_NAME)\n    model = AutoModel.from_pretrained(CFG.MODEL_NAME)\n    \n    # Chuyển sang Half Precision (FP16) để giảm 50% VRAM và tăng tốc\n    if CFG.DEVICE.type == 'cuda':\n        model = model.half()\n    \n    model.to(CFG.DEVICE)\n    model.eval()\n    return tokenizer, model\n\ndef extract_embeddings(fasta_path, prefix, tokenizer, model):\n    if not os.path.exists(fasta_path):\n        print(f\"Skipping {prefix}: File not found at {fasta_path}\")\n        return\n\n    print(f\"\\n--- Processing {prefix} Data ---\")\n    \n    # 1. Đọc dữ liệu từ FASTA\n    ids = []\n    seqs = []\n    for record in SeqIO.parse(fasta_path, \"fasta\"):\n        ids.append(str(record.id))\n        # Cắt sơ bộ sequence để tránh tốn RAM khi tạo list (model sẽ cắt chính xác sau)\n        seqs.append(str(record.seq)[:CFG.MAX_LEN + 50])\n        \n    n_samples = len(ids)\n    print(f\"Loaded {n_samples} sequences.\")\n\n    # 2. Sắp xếp theo độ dài (Smart Batching)\n    # Gom các chuỗi có độ dài tương đương vào 1 batch để giảm padding -> Chạy nhanh hơn 30%\n    sorted_indices = np.argsort([len(s) for s in seqs])[::-1] # Dài nhất xếp trước\n    \n    # 3. Khởi tạo mảng kết quả\n    # Output dim của ESM-2 650M là 1280\n    embed_dim = 1280 \n    embeddings_matrix = np.zeros((n_samples, embed_dim), dtype=np.float16)\n    \n    # 4. Vòng lặp Extract\n    with torch.inference_mode(): # Tắt gradient hoàn toàn\n        for i in tqdm(range(0, n_samples, CFG.BATCH_SIZE), desc=f\"Extracting {prefix}\"):\n            # Lấy batch index đã sort\n            batch_idx = sorted_indices[i : i + CFG.BATCH_SIZE]\n            batch_seqs = [seqs[idx] for idx in batch_idx]\n            \n            # Tokenize\n            inputs = tokenizer(\n                batch_seqs, \n                padding=True, \n                truncation=True, \n                max_length=CFG.MAX_LEN, \n                return_tensors=\"pt\"\n            )\n            \n            # Move to GPU\n            input_ids = inputs['input_ids'].to(CFG.DEVICE)\n            attention_mask = inputs['attention_mask'].to(CFG.DEVICE)\n            \n            # Forward Pass\n            output = model(input_ids=input_ids, attention_mask=attention_mask)\n            last_hidden_state = output.last_hidden_state # (Batch, Seq_Len, 1280)\n            \n            # Mean Pooling (Chỉ lấy trung bình các token thật, bỏ padding)\n            # Chú ý: convert mask sang float để nhân\n            mask_expanded = attention_mask.unsqueeze(-1).expand(last_hidden_state.size()).float()\n            \n            sum_embeddings = torch.sum(last_hidden_state * mask_expanded, 1)\n            sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)\n            \n            # Kết quả embedding cho batch này\n            batch_embeddings = (sum_embeddings / sum_mask).cpu().numpy().astype(np.float16)\n            \n            # Gán vào matrix tổng theo đúng vị trí index gốc\n            embeddings_matrix[batch_idx] = batch_embeddings\n            \n            # Dọn dẹp VRAM định kỳ\n            if i % 100 == 0:\n                del input_ids, attention_mask, output, last_hidden_state\n                torch.cuda.empty_cache()\n\n    # 5. Lưu File\n    out_emb_path = os.path.join(CFG.OUTPUT_DIR, f\"{prefix}_embeddings_esm2.npy\")\n    out_ids_path = os.path.join(CFG.OUTPUT_DIR, f\"{prefix}_ids_esm2.npy\")\n    \n    # Convert sang float32 khi save nếu bạn muốn an toàn tuyệt đối cho bước train, \n    # nhưng float16 cũng train tốt và nhẹ hơn. Ở đây tôi để float32 cho chuẩn.\n    np.save(out_emb_path, embeddings_matrix.astype(np.float32))\n    np.save(out_ids_path, np.array(ids))\n    \n    print(f\"Saved Embeddings: {out_emb_path} {embeddings_matrix.shape}\")\n    print(f\"Saved IDs: {out_ids_path}\")\n    \n    # Dọn RAM\n    del embeddings_matrix, ids, seqs, sorted_indices\n    gc.collect()\n\n# ==========================================\n# 3. THỰC THI (MAIN)\n# ==========================================\n\nif __name__ == \"__main__\":\n    # Setup\n    os.makedirs(CFG.OUTPUT_DIR, exist_ok=True)\n    \n    # Load Model (Load 1 lần dùng cho cả 2 file)\n    tokenizer, model = load_model()\n    \n    # Chạy Train\n    extract_embeddings(CFG.TRAIN_FASTA, \"train\", tokenizer, model)\n    \n    # Chạy Test (Nếu có)\n    extract_embeddings(CFG.TEST_FASTA, \"test\", tokenizer, model)\n    \n    print(\"\\nALL DONE! Embeddings are ready in Output section.\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-12-09T16:18:12.243946Z","iopub.execute_input":"2025-12-09T16:18:12.244479Z"}},"outputs":[],"execution_count":null}]}