{"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":31154,"isInternetEnabled":false,"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,"execution":{"iopub.status.busy":"2025-10-29T19:58:51.777451Z","iopub.execute_input":"2025-10-29T19:58:51.778110Z","iopub.status.idle":"2025-10-29T19:58:51.788825Z","shell.execute_reply.started":"2025-10-29T19:58:51.778086Z","shell.execute_reply":"2025-10-29T19:58:51.788117Z"}},"outputs":[],"execution_count":null},{"cell_type":"code","source":"import pandas as pd\nimport numpy as np\nfrom sklearn.preprocessing import MultiLabelBinarizer\nfrom sklearn.ensemble import RandomForestClassifier, GradientBoostingClassifier\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.model_selection import train_test_split\nfrom xgboost import XGBClassifier\nimport gc\nimport warnings\nwarnings.filterwarnings('ignore')\n\ntry:\n    import torch\n    from transformers import AutoTokenizer, AutoModel\n    HAS_ESM = True\nexcept Exception as e:\n    HAS_ESM = False\n    print(f\"Warning: ESM2 not available ({e}), using basic features\")\n    torch = None\n\nclass CAFA6ImprovedPredictor:\n    \"\"\"Improved CAFA-6 solution with ensemble & optimization\"\"\"\n    \n    def __init__(self):\n        self.models = {}\n        self.mlb = MultiLabelBinarizer()\n        self.top_go_terms = None\n        self.embedding_dim = None\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\") if HAS_ESM and torch is not None else None\n        self.scaler = None\n        \n    def load_fasta(self, fasta_path):\n        \"\"\"Load FASTA sequences\"\"\"\n        sequences = {}\n        try:\n            with open(fasta_path, 'r') as f:\n                current_id = None\n                current_seq = []\n                \n                for line in f:\n                    line = line.strip()\n                    if line.startswith('>'):\n                        if current_id:\n                            sequences[current_id] = ''.join(current_seq)\n                        current_id = line[1:].split()[0]\n                        current_seq = []\n                    elif current_id:\n                        current_seq.append(line)\n                \n                if current_id:\n                    sequences[current_id] = ''.join(current_seq)\n        except Exception as e:\n            print(f\"Error loading FASTA: {e}\")\n        \n        return sequences\n    \n    def load_esm_model(self):\n        \"\"\"Load ESM2 model\"\"\"\n        if not HAS_ESM or self.device is None:\n            return None, None\n        \n        try:\n            model_name = \"facebook/esm2_t12_35M_UR50D\"  # Better model (was t6_8M)\n            tokenizer = AutoTokenizer.from_pretrained(model_name)\n            model = AutoModel.from_pretrained(model_name).to(self.device).eval()\n            return tokenizer, model\n        except:\n            return None, None\n    \n    def get_batch_embeddings(self, sequences, tokenizer, model, batch_size=16):\n        \"\"\"Get embeddings with improved pooling\"\"\"\n        if tokenizer is None or model is None:\n            return np.array([self.get_advanced_features(seq) for seq in sequences])\n        \n        embeddings = []\n        \n        for i in range(0, len(sequences), batch_size):\n            batch_seqs = sequences[i:i+batch_size]\n            \n            try:\n                inputs = tokenizer(batch_seqs, return_tensors=\"pt\", padding=True, \n                                 truncation=True, max_length=1022).to(self.device)\n                \n                with torch.no_grad():\n                    outputs = model(**inputs)\n                    # Multi-pooling: mean + last token\n                    last_hidden = outputs.last_hidden_state\n                    mean_pool = last_hidden.mean(dim=1)\n                    cls_token = last_hidden[:, 0, :]\n                    # Concatenate different pooling strategies\n                    combined = torch.cat([mean_pool, cls_token], dim=1)\n                    embeddings.extend(combined.cpu().numpy())\n            except:\n                for seq in batch_seqs:\n                    embeddings.append(self.get_advanced_features(seq))\n        \n        return np.array(embeddings, dtype=np.float32)\n    \n    def get_advanced_features(self, sequence):\n        \"\"\"Enhanced sequence features\"\"\"\n        seq = str(sequence).upper()\n        length = max(len(seq), 1)\n        \n        features = []\n        \n        # Amino acid composition (20 standard)\n        aa_list = 'ACDEFGHIKLMNPQRSTVWY'\n        for aa in aa_list:\n            features.append(seq.count(aa) / length)\n        \n        # Chemical properties\n        hydrophobic = sum(seq.count(aa) for aa in 'AILMFVP') / length\n        charged = (sum(seq.count(aa) for aa in 'KR') - sum(seq.count(aa) for aa in 'DE')) / length\n        aromatic = (seq.count('F') + seq.count('W') + seq.count('Y')) / length\n        polar = sum(seq.count(aa) for aa in 'STNQ') / length\n        \n        # Di-peptide frequencies (sample)\n        dipeptides = ['AA', 'AC', 'AE', 'AG', 'AL']\n        for dp in dipeptides:\n            features.append(seq.count(dp) / max(length - 1, 1))\n        \n        features.extend([hydrophobic, charged, aromatic, polar, np.log1p(length)])\n        \n        return np.array(features, dtype=np.float32)\n    \n    def load_data(self, train_seq, train_taxon, train_terms, test_seq, test_taxon):\n        \"\"\"Load data\"\"\"\n        print(\"=\"*70)\n        print(\"LOADING DATA\")\n        print(\"=\"*70)\n        \n        print(\"\\n1. Loading sequences...\")\n        self.train_sequences = self.load_fasta(train_seq)\n        self.test_sequences = self.load_fasta(test_seq)\n        print(f\"   ✓ Train: {len(self.train_sequences)} | Test: {len(self.test_sequences)}\")\n        \n        print(\"\\n2. Loading taxonomy...\")\n        self.train_taxon = pd.read_csv(train_taxon, sep='\\t', header=None, \n                                       names=['protein_id', 'taxon_id'], dtype=str)\n        \n        print(\"\\n3. Loading annotations...\")\n        train_terms_df = pd.read_csv(train_terms, sep='\\t', header=None, \n                                     names=['protein_id', 'go_id', 'aspect'], dtype=str)\n        \n        # IMPROVEMENT: Select top 300 GO terms instead of 200\n        go_counts = train_terms_df['go_id'].value_counts()\n        self.top_go_terms = go_counts.head(300).index.tolist()\n        \n        train_terms_df = train_terms_df[train_terms_df['go_id'].isin(self.top_go_terms)]\n        self.go_dict = train_terms_df.groupby('protein_id')['go_id'].apply(list).to_dict()\n        \n        print(f\"   ✓ GO terms: {len(self.top_go_terms)}\")\n        print(f\"   ✓ Annotations: {len(train_terms_df)}\\n\")\n        \n        del train_terms_df\n        gc.collect()\n    \n    def prepare_embeddings(self):\n        \"\"\"Extract embeddings\"\"\"\n        print(\"=\"*70)\n        print(\"PREPARING EMBEDDINGS\")\n        print(\"=\"*70)\n        \n        protein_ids = self.train_taxon['protein_id'].values\n        sequences = [self.train_sequences.get(str(pid), '') for pid in protein_ids]\n        \n        print(\"\\nLoading ESM2 model...\")\n        tokenizer, model = self.load_esm_model()\n        \n        print(\"Extracting train embeddings...\")\n        X = self.get_batch_embeddings(sequences, tokenizer, model, batch_size=16)\n        \n        self.embedding_dim = X.shape[1]\n        print(f\"✓ Train embeddings: {X.shape}\")\n        \n        # Normalize embeddings (IMPROVEMENT)\n        from sklearn.preprocessing import StandardScaler\n        self.scaler = StandardScaler()\n        X = self.scaler.fit_transform(X)\n        \n        print(\"Creating labels...\")\n        self.go_labels = {}\n        for go_idx, go_term in enumerate(self.top_go_terms):\n            labels = []\n            for pid in protein_ids:\n                has_go = 1 if go_term in self.go_dict.get(str(pid), []) else 0\n                labels.append(has_go)\n            self.go_labels[go_idx] = np.array(labels, dtype=np.uint8)\n        \n        print(f\"✓ Labels created for {len(self.go_labels)} GO terms\\n\")\n        \n        return X, protein_ids\n    \n    def prepare_test_embeddings(self):\n        \"\"\"Extract test embeddings\"\"\"\n        print(\"Extracting test embeddings...\")\n        \n        test_ids = list(self.test_sequences.keys())\n        sequences = [self.test_sequences.get(str(pid), '') for pid in test_ids]\n        \n        print(\"Loading ESM2 model...\")\n        tokenizer, model = self.load_esm_model()\n        \n        X_test = self.get_batch_embeddings(sequences, tokenizer, model, batch_size=16)\n        X_test = self.scaler.transform(X_test)\n        \n        print(f\"✓ Test embeddings: {X_test.shape}\\n\")\n        \n        return X_test, test_ids\n    \n    def train_ensemble_models(self, X, protein_ids):\n        \"\"\"Train ensemble of models (IMPROVEMENT)\"\"\"\n        print(\"=\"*70)\n        print(\"TRAINING ENSEMBLE MODELS\")\n        print(\"=\"*70)\n        \n        # IMPROVEMENT: Use 70% of data for better training\n        sample_size = int(0.7 * len(protein_ids))\n        sample_idx = np.random.choice(len(protein_ids), size=sample_size, replace=False)\n        X_train = X[sample_idx]\n        \n        print(f\"\\nTraining on {len(X_train)} samples with ensemble...\\n\")\n        \n        for go_idx, go_term in enumerate(self.top_go_terms):\n            if go_idx % 50 == 0:\n                print(f\"Training GO {go_idx}/{len(self.top_go_terms)}...\")\n            \n            y_train = self.go_labels[go_idx][sample_idx]\n            \n            if y_train.sum() < 2:\n                self.models[go_idx] = None\n                continue\n            \n            # IMPROVEMENT: Ensemble of 3 models with voting\n            models_list = []\n            \n            # Model 1: Logistic Regression\n            lr = LogisticRegression(max_iter=300, random_state=42, class_weight='balanced', n_jobs=-1)\n            models_list.append(('lr', lr))\n            \n            # Model 2: XGBoost (better for complex patterns)\n            if y_train.sum() > 10:\n                xgb = XGBClassifier(max_depth=5, learning_rate=0.1, n_estimators=100, \n                                   random_state=42, scale_pos_weight=len(y_train)/y_train.sum())\n                models_list.append(('xgb', xgb))\n            \n            # Model 3: Random Forest (robustness)\n            rf = RandomForestClassifier(n_estimators=50, max_depth=10, random_state=42, \n                                       class_weight='balanced', n_jobs=-1)\n            models_list.append(('rf', rf))\n            \n            try:\n                trained_models = {}\n                for name, model in models_list:\n                    model.fit(X_train, y_train)\n                    trained_models[name] = model\n                self.models[go_idx] = trained_models\n            except:\n                self.models[go_idx] = None\n        \n        print(\"✓ Training complete!\\n\")\n    \n    def predict_probabilities(self, X_test):\n        \"\"\"Generate ensemble predictions (IMPROVEMENT)\"\"\"\n        print(\"=\"*70)\n        print(\"GENERATING ENSEMBLE PREDICTIONS\")\n        print(\"=\"*70 + \"\\n\")\n        \n        n_samples = X_test.shape[0]\n        n_go_terms = len(self.top_go_terms)\n        \n        predictions = np.zeros((n_samples, n_go_terms), dtype=np.float32)\n        \n        for go_idx in range(n_go_terms):\n            if go_idx % 50 == 0:\n                print(f\"Predicting GO {go_idx}/{n_go_terms}...\")\n            \n            model_dict = self.models.get(go_idx)\n            \n            if model_dict is None:\n                predictions[:, go_idx] = 0.1\n            else:\n                try:\n                    # Average predictions from ensemble\n                    ensemble_probs = []\n                    for name, model in model_dict.items():\n                        proba = model.predict_proba(X_test)\n                        if proba.shape[1] == 2:\n                            ensemble_probs.append(proba[:, 1])\n                    \n                    if ensemble_probs:\n                        predictions[:, go_idx] = np.mean(ensemble_probs, axis=0)\n                    else:\n                        predictions[:, go_idx] = 0.1\n                except:\n                    predictions[:, go_idx] = 0.1\n        \n        print(\"✓ Predictions generated!\\n\")\n        return predictions\n    \n    def create_submission(self, predictions, test_ids):\n        \"\"\"Create submission (IMPROVEMENT: better threshold tuning)\"\"\"\n        print(\"=\"*70)\n        print(\"CREATING SUBMISSION\")\n        print(\"=\"*70 + \"\\n\")\n        \n        submission_data = []\n        \n        for i, protein_id in enumerate(test_ids):\n            pred_probs = predictions[i]\n            \n            # IMPROVEMENT: Dynamic threshold based on probability distribution\n            threshold = np.percentile(pred_probs, 75)\n            threshold = max(threshold, 0.2)  # Minimum threshold\n            \n            top_indices = np.argsort(pred_probs)[-15:][::-1]  # Top 15 instead of 10\n            go_terms = []\n            \n            for idx in top_indices:\n                if pred_probs[idx] > threshold:\n                    go_terms.append(self.top_go_terms[idx])\n            \n            if not go_terms:\n                go_terms = [self.top_go_terms[top_indices[0]]]\n            \n            submission_data.append({\n                'target_id': protein_id,\n                'predictions': ' '.join(go_terms)\n            })\n            \n            if (i + 1) % 50000 == 0:\n                print(f\"  Processed {i + 1}/{len(test_ids)} predictions...\")\n        \n        submission_df = pd.DataFrame(submission_data)\n        submission_df.to_csv('submission.tsv', sep='\\t', index=False, header=False)\n        \n        print(f\"\\n✓ Submission saved: submission.tsv\")\n        print(f\"  Shape: {submission_df.shape}\")\n        print(f\"  Sample:\\n{submission_df.head()}\")\n    \n    def run(self, train_seq, train_taxon, train_terms, test_seq, test_taxon):\n        \"\"\"Execute pipeline\"\"\"\n        print(\"\\n\" + \"╔\" + \"=\"*68 + \"╗\")\n        print(\"║\" + \" \"*10 + \"CAFA-6 IMPROVED SOLUTION (0.32+)\" + \" \"*25 + \"║\")\n        print(\"║\" + \" \"*12 + \"ESM2 + ENSEMBLE + OPTIMIZATIONS\" + \" \"*23 + \"║\")\n        print(\"╚\" + \"=\"*68 + \"╝\\n\")\n        \n        try:\n            self.load_data(train_seq, train_taxon, train_terms, test_seq, test_taxon)\n            X, protein_ids = self.prepare_embeddings()\n            X_test, test_ids = self.prepare_test_embeddings()\n            self.train_ensemble_models(X, protein_ids)\n            predictions = self.predict_probabilities(X_test)\n            self.create_submission(predictions, test_ids)\n            \n            print(\"╔\" + \"=\"*68 + \"╗\")\n            print(\"║\" + \" \"*15 + \"✓ SUCCESS - READY FOR SUBMISSION!\" + \" \"*19 + \"║\")\n            print(\"╚\" + \"=\"*68 + \"╝\\n\")\n            \n        except Exception as e:\n            print(f\"\\n❌ Error: {e}\")\n            import traceback\n            traceback.print_exc()\n\n\nif __name__ == \"__main__\":\n    predictor = CAFA6ImprovedPredictor()\n    \n    predictor.run(\n        train_seq='/kaggle/input/cafa-6-protein-function-prediction/Train/train_sequences.fasta',\n        train_taxon='/kaggle/input/cafa-6-protein-function-prediction/Train/train_taxonomy.tsv',\n        train_terms='/kaggle/input/cafa-6-protein-function-prediction/Train/train_terms.tsv',\n        test_seq='/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset.fasta',\n        test_taxon='/kaggle/input/cafa-6-protein-function-prediction/Test/testsuperset-taxon-list.tsv'\n    )","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-10-29T20:05:49.366417Z","iopub.execute_input":"2025-10-29T20:05:49.367050Z","iopub.status.idle":"2025-10-29T20:11:12.088187Z","shell.execute_reply.started":"2025-10-29T20:05:49.367024Z","shell.execute_reply":"2025-10-29T20:11:12.087310Z"}},"outputs":[],"execution_count":null}]}