{"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":[{"sourceType":"competition","sourceId":4117,"databundleVersionId":46665}],"dockerImageVersionId":31328,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"import os, pandas as pd, numpy as np, subprocess, re, time, torch, gc\nimport torch.nn as nn\nimport torch.nn.functional as F\nimport torch.optim as optim\nfrom PIL import Image\nfrom collections import Counter\nfrom torchvision import transforms, models\nfrom torch.utils.data import DataLoader, TensorDataset\n\n# --- 1. CẤU HÌNH HỆ THỐNG ---\nTRAIN_7Z = '/kaggle/input/competitions/malware-classification/train.7z'\nTEST_7Z = '/kaggle/input/competitions/malware-classification/test.7z'\nTRAIN_LABELS = '/kaggle/input/competitions/malware-classification/trainLabels.csv'\nTEMP_DIR = '/kaggle/working/temp_data'\nSUBMISSION_FILE = 'submission_final.csv'\nos.makedirs(TEMP_DIR, exist_ok=True)\n\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\nopcodes = ['jmp', 'mov', 'retf', 'push', 'pop', 'add', 'sub', 'mul', 'div', 'and', 'or', 'xor', 'cmp', 'test', 'call']\nkeywords = ['db', 'dw', 'dd', 'offset', 'api', '.dll', 'std::']\nn_asm_features = len(opcodes) + len(keywords)\n\n# --- 2. HÀM TRÍCH XUẤT ĐẶC TRƯNG AN TOÀN ---\ndef extract_features_safe(fid):\n    b_path = os.path.join(TEMP_DIR, f\"{fid}.bytes\")\n    a_path = os.path.join(TEMP_DIR, f\"{fid}.asm\")\n    img_t = torch.zeros((1, 1, 128, 128))\n    asm_v = [0 for _ in range(n_asm_features)] \n    \n    if os.path.exists(b_path):\n        try:\n            with open(b_path, 'r') as f:\n                content = f.read(1000000)\n                hex_list = [int(x, 16) for x in content.split() if len(x)==2 and x!='??']\n            if hex_list:\n                d = np.array(hex_list, dtype=np.uint8)\n                side = int(len(d)**0.5)\n                if side > 10:\n                    img = Image.fromarray(d[:side*side].reshape((side, side))).convert('L').resize((128, 128))\n                    img_t = transforms.ToTensor()(img).unsqueeze(0)\n        except: pass\n\n    if os.path.exists(a_path):\n        try:\n            if os.path.getsize(a_path) < 150 * 1024 * 1024:\n                with open(a_path, 'r', encoding='latin-1') as f:\n                    content = f.read().lower()\n                    tokens = Counter(re.findall(r'\\b[a-z]+(?:::)?\\b', content))\n                    asm_v = [tokens.get(o, 0) for o in opcodes] + [tokens.get(k, 0) for k in keywords]\n        except: pass\n    return img_t, asm_v\n\n# --- 3. GIAI ĐOẠN 1: TRÍCH XUẤT DỮ LIỆU TRAIN ---\nprint(\">>> BẮT ĐẦU TRÍCH XUẤT TRAIN...\")\nlabels_df = pd.read_csv(TRAIN_LABELS)\ntrain_ids = labels_df['Id'].values\nX_asm, X_img, Y = [], [], []\n\nT_BATCH = 100 \nfor i in range(0, len(train_ids), T_BATCH):\n    batch = train_ids[i : i+T_BATCH]\n    with open('train_list.txt', 'w') as f:\n        for fid in batch: f.write(f\"train/{fid}.bytes\\ntrain/{fid}.asm\\n\")\n    \n    subprocess.run(f\"7z e {TRAIN_7Z} -o{TEMP_DIR} -i@train_list.txt -y > /dev/null\", shell=True)\n    \n    for fid in batch:\n        img, asm = extract_features_safe(fid)\n        X_img.append(img)\n        X_asm.append(asm)\n        # --- SỬA LỖI TẠI ĐÂY ---\n        # Thêm để lấy giá trị số nguyên đơn lẻ, không lấy mảng\n        label_val = int(labels_df[labels_df['Id']==fid]['Class'].values) - 1\n        Y.append(label_val)\n    \n    subprocess.run(f\"rm -rf {TEMP_DIR}/*\", shell=True)\n    if (i+T_BATCH) % 500 == 0: print(f\"Đã xử lý: {i+T_BATCH} file\")\n\n# Chuẩn bị Tensor (Sử dụng np.array trước để tăng tốc và fix dimension)\nX_asm_np = np.array(X_asm, dtype=np.float32)\nX_img_tensor = torch.cat(X_img, dim=0) \nY_np = np.array(Y).astype(np.int64) # Đã là mảng 1D nhờ bước sửa trên\nY_tensor = torch.tensor(Y_np, dtype=torch.long)\n\n# --- 4. GIAI ĐOẠN 2: HUẤN LUYỆN MODEL ---\nprint(\">>> ĐANG LUYỆN XGBOOST...\")\nimport xgboost as xgb\n\n# SỬA LỖI TẠI ĐÂY\nxgb_model = xgb.XGBClassifier(\n    tree_method='hist', \n    device='cuda' if torch.cuda.is_available() else 'cpu', \n    n_estimators=300\n)\n\nxgb_model.fit(X_asm_np, Y_np)\n# Lưu model ngay sau khi train xong\nxgb_model.save_model('malware_xgb_v2.json')\n\nprint(\">>> ĐANG LUYỆN RESNET...\")\nmodel_resnet = models.resnet18(weights=models.ResNet18_Weights.DEFAULT)\nmodel_resnet.conv1 = nn.Conv2d(1, 64, kernel_size=7, stride=2, padding=3, bias=False)\nmodel_resnet.fc = nn.Linear(512, 9)\nmodel_resnet.to(device)\n\ntrain_loader = DataLoader(TensorDataset(X_img_tensor, Y_tensor), batch_size=64, shuffle=True)\noptimizer = optim.Adam(model_resnet.parameters(), lr=0.001)\ncriterion = nn.CrossEntropyLoss()\n\nmodel_resnet.train()\nfor epoch in range(10):\n    for inputs, labels in train_loader:\n        inputs, labels = inputs.to(device), labels.to(device)\n        optimizer.zero_grad()\n        # Chạy forward và tính loss - labels bây giờ chuẩn 1D tensor\n        outputs = model_resnet(inputs)\n        loss = criterion(outputs, labels)\n        loss.backward()\n        optimizer.step()\n    print(f\"Epoch {epoch+1} hoàn tất.\")\n\ntorch.save(model_resnet.state_dict(), 'malware_resnet_v2.pth')\nprint(\"✅ Đã lưu cả 2 model thành công.\")\n\n# Giải phóng bộ nhớ\ndel X_img, X_asm, X_img_tensor, Y_tensor\ngc.collect()\ntorch.cuda.empty_cache()\n\n# --- 5. GIAI ĐOẠN 3: DỰ ĐOÁN TẬP TEST ---\nprint(\">>> ĐANG DỰ ĐOÁN TẬP TEST...\")\nraw_test = subprocess.check_output(f\"7z l {TEST_7Z} | grep .bytes | awk '{{print $NF}}'\", shell=True).decode().splitlines()\ntest_ids = [os.path.basename(f).replace('.bytes', '') for f in raw_test if len(os.path.basename(f)) > 15]\n\nif os.path.exists(SUBMISSION_FILE): os.remove(SUBMISSION_FILE)\n\nmodel_resnet.eval()\nB_SIZE = 100\nfor i in range(0, len(test_ids), B_SIZE):\n    batch = test_ids[i : i+B_SIZE]\n    with open('test_batch.txt', 'w') as f:\n        for fid in batch: f.write(f\"test/{fid}.bytes\\ntest/{fid}.asm\\n\")\n    \n    subprocess.run(f\"7z e {TEST_7Z} -o{TEMP_DIR} -i@test_batch.txt -y > /dev/null\", shell=True)\n    \n    batch_probs = []\n    for fid in batch:\n        img_t, asm_f = extract_features_safe(fid)\n        with torch.no_grad():\n            p_resnet = F.softmax(model_resnet(img_t.to(device)), dim=1).cpu().numpy().flatten().tolist()\n        \n        p_xgb_raw = xgb_model.predict_proba(np.array([asm_f])).flatten().tolist()\n        \n        final_p_xgb = [0.0] * 9\n        for idx in range(min(len(p_xgb_raw), 9)):\n            final_p_xgb[idx] = p_xgb_raw[idx]\n            \n        combined = []\n        for r, x in zip(p_resnet, final_p_xgb):\n            combined.append(0.4 * r + 0.6 * x)\n        \n        batch_probs.append([fid] + combined)\n    \n    cols = ['Id'] + [f'Prediction{j}' for j in range(1, 10)] # Tên cột theo chuẩn Kaggle thường là Class_1..9\n    df_batch = pd.DataFrame(batch_probs, columns=cols)\n    df_batch.to_csv(SUBMISSION_FILE, mode='a', header=not os.path.exists(SUBMISSION_FILE), index=False)\n    \n    subprocess.run(f\"rm -rf {TEMP_DIR}/*\", shell=True)\n    if (i+B_SIZE) % 500 == 0:\n        print(f\"Tiến độ Test: {min(i+B_SIZE, len(test_ids))}/{len(test_ids)}\")\n        gc.collect()\n        torch.cuda.empty_cache()\n\nprint(f\"🏁 Xong! File: {SUBMISSION_FILE}\")","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-04-22T07:20:12.101709Z","iopub.execute_input":"2026-04-22T07:20:12.101987Z","execution_failed":"2026-04-22T07:20:11.976Z"}},"outputs":[],"execution_count":null}]}