{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"codemirror_mode":{"name":"ipython","version":3},"file_extension":".py","mimetype":"text/x-python","name":"python","version":"3.10.12"},"kaggle":{"accelerator":"gpu","dataSources":[{"sourceId":106680,"databundleVersionId":13374319,"sourceType":"competition"}],"dockerImageVersionId":31193,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":4,"nbformat":4,"cells":[{"id":"1a8507ed-6b52-4ff7-aca1-549ec4f52a90","cell_type":"markdown","source":"# AIRR-ML-25: Adaptive Immune Profilling \n\n## **Author**: Nayananshu Garai ","metadata":{}},{"id":"b23d1f2c-f42c-4038-9426-09342d8ad057","cell_type":"code","source":"\nimport os, sys, gc, warnings\nimport numpy as np\nimport pandas as pd\nfrom collections import Counter\nfrom tqdm.auto import tqdm\nfrom scipy import stats\n\nimport torch\nfrom transformers import AutoTokenizer, AutoModel\n\nfrom sklearn.model_selection import StratifiedKFold\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.metrics import roc_auc_score\nfrom sklearn.feature_selection import SelectKBest, f_classif\n\nimport xgboost as xgb\nimport lightgbm as lgb\n\nwarnings.filterwarnings('ignore')\nnp.random.seed(42)\ntorch.manual_seed(42)\n\ndevice = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\nprint(f\"✅ Device: {device}\")\nprint(f\"PyTorch: {torch.__version__}\")\nif torch.cuda.is_available():\n    print(f\"GPU: {torch.cuda.get_device_name(0)}\")\n    print(f\"Memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB\")\n    # Optimize GPU settings\n    torch.backends.cudnn.benchmark = True\n    torch.backends.cuda.matmul.allow_tf32 = True\nelse:\n    print(\"⚠️ GPU not detected! Enable in Settings\")","metadata":{},"outputs":[],"execution_count":null},{"id":"5160f715-b572-44c5-ae99-bcfd72a66eaf","cell_type":"code","source":"# Load data\nDATA_DIR = \"/kaggle/input/adaptive-immune-profiling-challenge-2025\"\nTRAIN_DIR = os.path.join(DATA_DIR, \"train_datasets\", \"train_datasets\")\nTEST_DIR = os.path.join(DATA_DIR, \"test_datasets\", \"test_datasets\")\n\ntrain_datasets = sorted([d for d in os.listdir(TRAIN_DIR) if d.startswith('train_dataset_')])\ntest_datasets = sorted([d for d in os.listdir(TEST_DIR) if d.startswith('test_dataset_')])\n\ntraining_metadata = {}\ntest_metadata = {}\n\nfor train_ds in train_datasets:\n    meta = pd.read_csv(os.path.join(TRAIN_DIR, train_ds, 'metadata.csv'))\n    training_metadata[train_ds] = meta\n\nfor test_ds in test_datasets:\n    meta_path = os.path.join(TEST_DIR, test_ds, 'metadata.csv')\n    if os.path.exists(meta_path):\n        meta = pd.read_csv(meta_path)\n    else:\n        tsv_files = [f for f in os.listdir(os.path.join(TEST_DIR, test_ds)) if f.endswith('.tsv')]\n        meta = pd.DataFrame([{'repertoire_id': f.replace('.tsv', ''), 'filename': f} for f in tsv_files])\n    test_metadata[test_ds] = meta\n\nprint(f\"✅ {len(train_datasets)} train + {len(test_datasets)} test datasets\")\nprint(f\"📊 {sum(len(m) for m in training_metadata.values())} train repertoires\")\nprint(f\"📊 {sum(len(m) for m in test_metadata.values())} test repertoires\")","metadata":{},"outputs":[],"execution_count":null},{"id":"bd776311-c219-43a2-a2ce-61ba8e48e590","cell_type":"code","source":"# Load ESM-2\nprint(\"⏳ Loading ESM-2 (2-3 min)...\")\n\nMODEL_NAME = \"facebook/esm2_t30_150M_UR50D\"\ntokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)\nmodel = AutoModel.from_pretrained(MODEL_NAME)\nmodel = model.to(device)\nmodel.eval()\n\n# Enable mixed precision for speed\nif device.type == 'cuda':\n    model = model.half()  # FP16 for 2x speed\n\nprint(f\"✅ ESM-2 loaded (150M params, FP16 mode)\")","metadata":{},"outputs":[],"execution_count":null},{"id":"d1a70fd0-e6ed-4a80-a398-e8ddcab84824","cell_type":"code","source":"# Feature extraction functions (OPTIMIZED)\n\n@torch.no_grad()\ndef get_esm_embeddings(sequences, batch_size=64):  # Increased batch size\n    \"\"\"Fast ESM-2 extraction with FP16.\"\"\"\n    all_embeddings = []\n    \n    for i in range(0, len(sequences), batch_size):\n        batch = sequences[i:i+batch_size]\n        inputs = tokenizer(batch, return_tensors='pt', padding=True, \n                          truncation=True, max_length=256)  # Reduced from 512\n        inputs = {k: v.to(device) for k, v in inputs.items()}\n        \n        outputs = model(**inputs)\n        embeddings = outputs.last_hidden_state[:, 0, :].cpu().float().numpy()\n        all_embeddings.append(embeddings)\n    \n    all_embeddings = np.vstack(all_embeddings)\n    \n    # Triple pooling\n    emb_mean = np.mean(all_embeddings, axis=0)\n    emb_max = np.max(all_embeddings, axis=0)\n    emb_std = np.std(all_embeddings, axis=0)\n    \n    return np.concatenate([emb_mean, emb_max, emb_std])\n\ndef extract_basic_features(df):\n    \"\"\"Essential statistical features only.\"\"\"\n    features = {}\n    \n    # Core stats\n    features['n_sequences'] = len(df)\n    features['n_unique'] = df['junction_aa'].nunique()\n    features['diversity'] = features['n_unique'] / features['n_sequences']\n    \n    # Length\n    lengths = df['junction_aa'].str.len()\n    features['length_mean'] = lengths.mean()\n    features['length_std'] = lengths.std()\n    features['length_median'] = lengths.median()\n    \n    # Gene usage\n    features['n_v_genes'] = df['v_call'].nunique()\n    features['n_j_genes'] = df['j_call'].nunique()\n    \n    v_counts = df['v_call'].value_counts()\n    j_counts = df['j_call'].value_counts()\n    \n    features['v_entropy'] = stats.entropy(v_counts.values)\n    features['j_entropy'] = stats.entropy(j_counts.values)\n    features['v_gini'] = 1 - sum((v_counts / len(df)) ** 2)\n    features['j_gini'] = 1 - sum((j_counts / len(df)) ** 2)\n    \n    # V-J pairing\n    vj_pairs = df.groupby(['v_call', 'j_call']).size()\n    features['n_vj_pairs'] = len(vj_pairs)\n    features['vj_entropy'] = stats.entropy(vj_pairs.values)\n    \n    return features\n\nprint(\"✅ Functions ready\")","metadata":{},"outputs":[],"execution_count":null},{"id":"97de6392-b8b1-42c6-81f3-cf214858afc5","cell_type":"code","source":"# Extract TRAINING features (FAST MODE)\nprint(f\"\\n{'='*70}\")\nprint(\"⚡ TRAINING FEATURES (Speed Mode: 50 seqs/file)\")\nprint(f\"{'='*70}\\n\")\n\ntraining_features = {}\n\nfor train_ds in train_datasets:\n    print(f\"📊 {train_ds}\")\n    dataset_path = os.path.join(TRAIN_DIR, train_ds)\n    meta = training_metadata[train_ds]\n    \n    features_list = []\n    \n    for idx, row in tqdm(meta.iterrows(), total=len(meta), desc=\"  \", leave=False):\n        try:\n            file_path = os.path.join(dataset_path, row['filename'])\n            df = pd.read_csv(file_path, sep='\\t')\n            \n            # Speed optimization: Only 50 sequences\n            sample_size = min(50, len(df))  # Was 300!\n            sampled_df = df.sample(n=sample_size, random_state=42)\n            sequences = sampled_df['junction_aa'].tolist()\n            \n            # ESM embeddings (960D)\n            esm_features = get_esm_embeddings(sequences, batch_size=64)\n            \n            # Basic features (computed on ALL sequences for accuracy)\n            basic_features = extract_basic_features(df)\n            \n            # Combine\n            all_features = {'repertoire_id': row['repertoire_id']}\n            for i, val in enumerate(esm_features):\n                all_features[f'esm_{i}'] = val\n            all_features.update(basic_features)\n            \n            if 'label_positive' in row:\n                all_features['label_positive'] = row['label_positive']\n            \n            features_list.append(all_features)\n            \n        except Exception as e:\n            continue\n    \n    training_features[train_ds] = pd.DataFrame(features_list)\n    print(f\"  ✅ {len(training_features[train_ds])} reps | {len(training_features[train_ds].columns)} features\\n\")\n    \n    torch.cuda.empty_cache()\n    gc.collect()\n\nprint(f\"{'='*70}\")\nprint(\"✅ TRAINING DONE\")\nprint(f\"{'='*70}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"22922bfb-964f-4cfd-955e-9fa34321b2e5","cell_type":"code","source":"# Extract TEST features (FAST MODE)\nprint(f\"\\n{'='*70}\")\nprint(\"⚡ TEST FEATURES (Speed Mode: 50 seqs/file)\")\nprint(f\"{'='*70}\\n\")\n\ntest_features = {}\n\nfor test_ds in test_datasets:\n    print(f\"📊 {test_ds}\")\n    dataset_path = os.path.join(TEST_DIR, test_ds)\n    meta = test_metadata[test_ds]\n    \n    features_list = []\n    \n    for idx, row in tqdm(meta.iterrows(), total=len(meta), desc=\"  \", leave=False):\n        try:\n            file_path = os.path.join(dataset_path, row['filename'])\n            df = pd.read_csv(file_path, sep='\\t')\n            \n            sample_size = min(50, len(df))\n            sampled_df = df.sample(n=sample_size, random_state=42)\n            sequences = sampled_df['junction_aa'].tolist()\n            \n            esm_features = get_esm_embeddings(sequences, batch_size=64)\n            basic_features = extract_basic_features(df)\n            \n            all_features = {'repertoire_id': row['repertoire_id']}\n            for i, val in enumerate(esm_features):\n                all_features[f'esm_{i}'] = val\n            all_features.update(basic_features)\n            \n            features_list.append(all_features)\n            \n        except Exception as e:\n            continue\n    \n    test_features[test_ds] = pd.DataFrame(features_list)\n    print(f\"  ✅ {len(test_features[test_ds])} reps\\n\")\n    \n    torch.cuda.empty_cache()\n    gc.collect()\n\nprint(f\"{'='*70}\")\nprint(\"✅ TEST DONE\")\nprint(f\"{'='*70}\")\n\n# Free GPU\ndel model, tokenizer\ntorch.cuda.empty_cache()\ngc.collect()\nprint(\"\\n🗑️ GPU memory freed\")","metadata":{},"outputs":[],"execution_count":null},{"id":"5fea185c-e408-40a1-9518-23fedd879cb1","cell_type":"code","source":"# Train models (FAST: 2 best models only)\nprint(f\"\\n{'='*70}\")\nprint(\"🚀 FAST TRAINING (XGBoost + LightGBM)\")\nprint(f\"{'='*70}\")\n\ntrained_models = {}\nall_predictions = {}\n\nfor train_ds in train_datasets:\n    print(f\"\\n📊 {train_ds}\")\n    \n    train_df = training_features[train_ds].copy()\n    feature_cols = [c for c in train_df.columns if c not in ['repertoire_id', 'label_positive']]\n    X = train_df[feature_cols].fillna(0)\n    y = train_df['label_positive'].astype(int)\n    \n    print(f\"  Samples: {len(X)} | Features: {len(feature_cols)}\")\n    \n    # Feature selection: Keep top 100 features for speed\n    if len(feature_cols) > 100:\n        selector = SelectKBest(f_classif, k=100)\n        X_selected = selector.fit_transform(X, y)\n        selected_features = X.columns[selector.get_support()].tolist()\n        X = pd.DataFrame(X_selected, columns=selected_features, index=X.index)\n        feature_cols = selected_features\n        print(f\"  Selected: {len(feature_cols)} best features\")\n    \n    # Scale\n    scaler = StandardScaler()\n    X_scaled = scaler.fit_transform(X)\n    X_scaled = pd.DataFrame(X_scaled, columns=feature_cols, index=X.index)\n    \n    # Class imbalance\n    pos_ratio = y.sum() / len(y)\n    is_imbalanced = pos_ratio < 0.3 or pos_ratio > 0.7\n    scale_pos_weight = (len(y) - y.sum()) / y.sum() if (is_imbalanced and y.sum() > 0) else 1\n    \n    print(f\"  Training 2 models...\")\n    \n    # XGBoost (BEST for embeddings)\n    xgb_model = xgb.XGBClassifier(\n        n_estimators=500, max_depth=6, learning_rate=0.02,\n        subsample=0.8, colsample_bytree=0.8,\n        scale_pos_weight=scale_pos_weight,\n        reg_alpha=0.3, reg_lambda=2.0,\n        random_state=42, n_jobs=-1, tree_method='gpu_hist',  # GPU acceleration!\n        eval_metric='logloss'\n    )\n    xgb_model.fit(X_scaled, y, verbose=False)\n    \n    # LightGBM (FAST)\n    lgb_model = lgb.LGBMClassifier(\n        n_estimators=500, max_depth=6, learning_rate=0.02,\n        subsample=0.8, colsample_bytree=0.8,\n        class_weight='balanced' if is_imbalanced else None,\n        reg_alpha=0.3, reg_lambda=2.0,\n        random_state=42, n_jobs=-1, device='gpu',  # GPU!\n        verbose=-1\n    )\n    lgb_model.fit(X_scaled, y)\n    \n    # Quick 3-fold CV\n    skf = StratifiedKFold(n_splits=3, shuffle=True, random_state=42)\n    cv_xgb, cv_lgb, cv_ensemble = [], [], []\n    \n    for train_idx, val_idx in skf.split(X_scaled, y):\n        X_tr, X_val = X_scaled.iloc[train_idx], X_scaled.iloc[val_idx]\n        y_tr, y_val = y.iloc[train_idx], y.iloc[val_idx]\n        \n        xgb_cv = xgb.XGBClassifier(n_estimators=300, max_depth=6, learning_rate=0.02,\n                                  scale_pos_weight=scale_pos_weight, random_state=42, \n                                  n_jobs=-1, tree_method='gpu_hist')\n        lgb_cv = lgb.LGBMClassifier(n_estimators=300, max_depth=6, learning_rate=0.02,\n                                   class_weight='balanced' if is_imbalanced else None,\n                                   random_state=42, n_jobs=-1, device='gpu', verbose=-1)\n        \n        xgb_cv.fit(X_tr, y_tr, verbose=False)\n        lgb_cv.fit(X_tr, y_tr)\n        \n        pred_xgb = xgb_cv.predict_proba(X_val)[:, 1]\n        pred_lgb = lgb_cv.predict_proba(X_val)[:, 1]\n        pred_ensemble = 0.6*pred_xgb + 0.4*pred_lgb  # Weighted\n        \n        cv_xgb.append(roc_auc_score(y_val, pred_xgb))\n        cv_lgb.append(roc_auc_score(y_val, pred_lgb))\n        cv_ensemble.append(roc_auc_score(y_val, pred_ensemble))\n    \n    print(f\"  CV Results (3-fold):\")\n    print(f\"    XGBoost:  {np.mean(cv_xgb):.4f} (±{np.std(cv_xgb):.4f})\")\n    print(f\"    LightGBM: {np.mean(cv_lgb):.4f} (±{np.std(cv_lgb):.4f})\")\n    print(f\"    Ensemble: {np.mean(cv_ensemble):.4f} (±{np.std(cv_ensemble):.4f})\")\n    \n    # Store\n    trained_models[train_ds] = {\n        'xgb': xgb_model,\n        'lgb': lgb_model,\n        'features': feature_cols,\n        'scaler': scaler\n    }\n    \n    # Predict test\n    test_matches = [ds for ds in test_datasets if train_ds.split('_')[-1] in ds]\n    \n    for test_ds in test_matches:\n        test_df = test_features[test_ds].copy()\n        \n        for col in feature_cols:\n            if col not in test_df.columns:\n                test_df[col] = 0\n        \n        X_test = test_df[feature_cols].fillna(0)\n        X_test_scaled = scaler.transform(X_test)\n        \n        pred_xgb = xgb_model.predict_proba(X_test_scaled)[:, 1]\n        pred_lgb = lgb_model.predict_proba(X_test_scaled)[:, 1]\n        pred_ensemble = 0.6*pred_xgb + 0.4*pred_lgb\n        \n        all_predictions[test_ds] = {\n            'repertoire_id': test_df['repertoire_id'].values,\n            'predictions': pred_ensemble\n        }\n        \n        print(f\"  {test_ds}: [{pred_ensemble.min():.3f}, {pred_ensemble.max():.3f}]\")\n    \n    gc.collect()\n\nprint(f\"\\n{'='*70}\")\nprint(\"✅ TRAINING COMPLETE\")\nprint(f\"{'='*70}\")","metadata":{},"outputs":[],"execution_count":null},{"id":"4db4072e-33b5-44fc-bf26-2ae6d35ec073","cell_type":"code","source":"# Create submission\nprint(\"\\n⏳ Creating submission...\\n\")\n\nsample_sub = pd.read_csv(\"/kaggle/input/adaptive-immune-profiling-challenge-2025/sample_submissions.csv\")\nsubmission_rows = []\n\n# Test predictions\nfor test_ds in test_datasets:\n    if test_ds in all_predictions:\n        pred_data = all_predictions[test_ds]\n        for rep_id, prob in zip(pred_data['repertoire_id'], pred_data['predictions']):\n            submission_rows.append({\n                'ID': rep_id,\n                'dataset': test_ds,\n                'label_positive_probability': float(prob),\n                'junction_aa': -999.0,\n                'v_call': -999.0,\n                'j_call': -999.0\n            })\n\n# Sequence rankings\nfor train_ds in train_datasets:\n    sample_seqs = sample_sub[sample_sub['dataset'] == train_ds]\n    if len(sample_seqs) > 0:\n        for idx, row in sample_seqs.iterrows():\n            if row['junction_aa'] != -999.0:\n                submission_rows.append({\n                    'ID': row['ID'],\n                    'dataset': train_ds,\n                    'label_positive_probability': -999.0,\n                    'junction_aa': row['junction_aa'],\n                    'v_call': row['v_call'],\n                    'j_call': row['j_call']\n                })\n\nsubmission_df = pd.DataFrame(submission_rows)\nsubmission_df.to_csv(\"submissions.csv\", index=False)\n\nprint(f\"{'='*70}\")\nprint(f\"✅ SUBMISSION READY\")\nprint(f\"{'='*70}\")\nprint(f\"File: submissions.csv\")\nprint(f\"Shape: {submission_df.shape}\")\nprint(f\"Size: {os.path.getsize('submissions.csv') / 1e6:.1f} MB\")\nprint(f\"\\n🚀 Ready to submit!\")","metadata":{},"outputs":[],"execution_count":null}]}