{"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":"none","dataSources":[{"sourceId":106680,"databundleVersionId":13374319,"sourceType":"competition"}],"dockerImageVersionId":31153,"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":false}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"markdown","source":"# Building a Baseline Model\n\nIn this notebook, we will build a simple baseline model for addressing the two tasks of the challenge: (a) predicting labels of examples, and (b) identifying top instances that best explain the label-prediction. We will adhere to the provided code template when building the model.\n\nThroughout the descriptions below, the language being used is deliberately generic to cater to a broader audience outside the immunology field. \n\nIn this example baseline model, we will use k-mer encoding - counting short subsequences of a defined size (here k=3) combined with L1-regularized logistic regression. This approach could be one intuitive baseline for a problem like this.","metadata":{}},{"cell_type":"code","source":"## imports used by the basic code template provided.\n\nimport os\nfrom tqdm import tqdm\nimport pandas as pd\nimport numpy as np\nimport torch\nimport glob\nimport sys\nimport argparse\nfrom collections import defaultdict\nfrom typing import Iterator, Tuple, Union, List\n\n## imports that are additionally used by this notebook\n\nfrom sklearn.linear_model import LogisticRegression\nfrom sklearn.preprocessing import StandardScaler\nfrom sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split\nfrom sklearn.metrics import roc_auc_score, balanced_accuracy_score\nfrom sklearn.pipeline import Pipeline\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T13:55:40.525949Z","iopub.execute_input":"2025-11-04T13:55:40.526245Z","iopub.status.idle":"2025-11-04T13:55:47.789321Z","shell.execute_reply.started":"2025-11-04T13:55:40.526219Z","shell.execute_reply":"2025-11-04T13:55:47.788310Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Utility functions\n\nThe code template provides some utility functions such as data loaders. In this example, I do not change anything from the chunk below, and just retain what was provided. One of the utility function `load_and_encode_kmers` reads in each example and encodes k-mer without consuming a lot of memory. We will use that later on when ecnoding data.","metadata":{}},{"cell_type":"code","source":"## some utility functions such as data loaders, etc.\n\ndef load_data_generator(data_dir: str, metadata_filename='metadata.csv') -> Iterator[\n    Union[Tuple[str, pd.DataFrame, bool], Tuple[str, pd.DataFrame]]]:\n    \"\"\"\n    A generator to load immune repertoire data.\n\n    This function operates in two modes:\n    1.  If metadata is found, it yields data based on the metadata file.\n    2.  If metadata is NOT found, it uses glob to find and yield all '.tsv'\n        files in the directory.\n\n    Args:\n        data_dir (str): The path to the directory containing the data.\n\n    Yields:\n        An iterator of tuples. The format depends on the mode:\n        - With metadata: (repertoire_id, pd.DataFrame, label_positive)\n        - Without metadata: (filename, pd.DataFrame)\n    \"\"\"\n    metadata_path = os.path.join(data_dir, metadata_filename)\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        for row in metadata_df.itertuples(index=False):\n            file_path = os.path.join(data_dir, row.filename)\n            try:\n                repertoire_df = pd.read_csv(file_path, sep='\\t')\n                yield row.repertoire_id, repertoire_df, row.label_positive\n            except FileNotFoundError:\n                print(f\"Warning: File '{row.filename}' listed in metadata not found. Skipping.\")\n                continue\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        tsv_files = glob.glob(search_pattern)\n        for file_path in sorted(tsv_files):\n            try:\n                filename = os.path.basename(file_path)\n                repertoire_df = pd.read_csv(file_path, sep='\\t')\n                yield filename, repertoire_df\n            except Exception as e:\n                print(f\"Warning: Could not read file '{file_path}'. Error: {e}. Skipping.\")\n                continue\n\n\ndef load_full_dataset(data_dir: str) -> pd.DataFrame:\n    \"\"\"\n    Loads all TSV files from a directory and concatenates them into a single DataFrame.\n\n    This function handles two scenarios:\n    1. If metadata.csv exists, it loads data based on the metadata and adds\n       'repertoire_id' and 'label_positive' columns.\n    2. If metadata.csv does not exist, it loads all .tsv files and adds\n       a 'filename' column as an identifier.\n\n    Args:\n        data_dir (str): The path to the data directory.\n\n    Returns:\n        pd.DataFrame: A single, concatenated DataFrame containing all the data.\n    \"\"\"\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n    df_list = []\n    data_loader = load_data_generator(data_dir=data_dir)\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        total_files = len(metadata_df)\n        for rep_id, data_df, label in tqdm(data_loader, total=total_files, desc=\"Loading files\"):\n            data_df['ID'] = rep_id\n            data_df['label_positive'] = label\n            df_list.append(data_df)\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        total_files = len(glob.glob(search_pattern))\n        for filename, data_df in tqdm(data_loader, total=total_files, desc=\"Loading files\"):\n            data_df['ID'] = os.path.basename(filename).replace(\".tsv\", \"\")\n            df_list.append(data_df)\n\n    if not df_list:\n        print(\"Warning: No data files were loaded.\")\n        return pd.DataFrame()\n\n    full_dataset_df = pd.concat(df_list, ignore_index=True)\n    return full_dataset_df\n\n\ndef load_and_encode_kmers(data_dir: str, k: int = 3) -> Tuple[pd.DataFrame, pd.DataFrame]:\n    \"\"\"\n    Loading and k-mer encoding of repertoire data.\n\n    Args:\n        data_dir: Path to data directory\n        k: K-mer length\n\n    Returns:\n        Tuple of (encoded_features_df, metadata_df)\n        metadata_df always contains 'ID', and 'label_positive' if available\n    \"\"\"\n    from collections import Counter\n\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n    data_loader = load_data_generator(data_dir=data_dir)\n\n    repertoire_features = []\n    metadata_records = []\n\n    search_pattern = os.path.join(data_dir, '*.tsv')\n    total_files = len(glob.glob(search_pattern))\n\n    for item in tqdm(data_loader, total=total_files, desc=f\"Encoding {k}-mers\"):\n        if os.path.exists(metadata_path):\n            rep_id, data_df, label = item\n        else:\n            filename, data_df = item\n            rep_id = os.path.basename(filename).replace(\".tsv\", \"\")\n            label = None\n\n        kmer_counts = Counter()\n        for seq in data_df['junction_aa'].dropna():\n            for i in range(len(seq) - k + 1):\n                kmer_counts[seq[i:i + k]] += 1\n\n        repertoire_features.append({\n            'ID': rep_id,\n            **kmer_counts\n        })\n\n        metadata_record = {'ID': rep_id}\n        if label is not None:\n            metadata_record['label_positive'] = label\n        metadata_records.append(metadata_record)\n\n        del data_df, kmer_counts\n\n    features_df = pd.DataFrame(repertoire_features).fillna(0).set_index('ID')\n    features_df.fillna(0)\n    metadata_df = pd.DataFrame(metadata_records)\n\n    return features_df, metadata_df\n\n\ndef save_tsv(df: pd.DataFrame, path: str):\n    os.makedirs(os.path.dirname(path), exist_ok=True)\n    df.to_csv(path, sep='\\t', index=False)\n\n\ndef get_repertoire_ids(data_dir: str) -> list:\n    \"\"\"\n    Retrieves repertoire IDs from the metadata file or filenames in the directory.\n\n    Args:\n        data_dir (str): The path to the data directory.\n\n    Returns:\n        list: A list of repertoire IDs.\n    \"\"\"\n    metadata_path = os.path.join(data_dir, 'metadata.csv')\n\n    if os.path.exists(metadata_path):\n        metadata_df = pd.read_csv(metadata_path)\n        repertoire_ids = metadata_df['repertoire_id'].tolist()\n    else:\n        search_pattern = os.path.join(data_dir, '*.tsv')\n        tsv_files = glob.glob(search_pattern)\n        repertoire_ids = [os.path.basename(f).replace('.tsv', '') for f in sorted(tsv_files)]\n\n    return repertoire_ids\n\n\ndef generate_random_top_sequences_df(n_seq: int = 50000) -> pd.DataFrame:\n    \"\"\"\n    Generates a random DataFrame simulating top important sequences.\n\n    Args:\n        n_seq (int): Number of sequences to generate.\n\n    Returns:\n        pd.DataFrame: A DataFrame with columns 'ID', 'dataset', 'junction_aa', 'v_call', 'j_call'.\n    \"\"\"\n    seqs = set()\n    while len(seqs) < n_seq:\n        seq = ''.join(np.random.choice(list('ACDEFGHIKLMNPQRSTVWY'), size=15))\n        seqs.add(seq)\n    data = {\n        'junction_aa': list(seqs),\n        'v_call': ['TRBV20-1'] * n_seq,\n        'j_call': ['TRBJ2-7'] * n_seq,\n        'importance_score': np.random.rand(n_seq)\n    }\n    return pd.DataFrame(data)\n\n\ndef validate_dirs_and_files(train_dir: str, test_dirs: List[str], out_dir: str) -> None:\n    assert os.path.isdir(train_dir), f\"Train directory `{train_dir}` does not exist.\"\n    train_tsvs = glob.glob(os.path.join(train_dir, \"*.tsv\"))\n    assert train_tsvs, f\"No .tsv files found in train directory `{train_dir}`.\"\n    metadata_path = os.path.join(train_dir, \"metadata.csv\")\n    assert os.path.isfile(metadata_path), f\"`metadata.csv` not found in train directory `{train_dir}`.\"\n\n    for test_dir in test_dirs:\n        assert os.path.isdir(test_dir), f\"Test directory `{test_dir}` does not exist.\"\n        test_tsvs = glob.glob(os.path.join(test_dir, \"*.tsv\"))\n        assert test_tsvs, f\"No .tsv files found in test directory `{test_dir}`.\"\n\n    try:\n        os.makedirs(out_dir, exist_ok=True)\n        test_file = os.path.join(out_dir, \"test_write_permission.tmp\")\n        with open(test_file, \"w\") as f:\n            f.write(\"test\")\n        os.remove(test_file)\n    except Exception as e:\n        print(f\"Failed to create or write to output directory `{out_dir}`: {e}\")\n        sys.exit(1)\n\n\ndef concatenate_output_files(out_dir: str) -> None:\n    \"\"\"\n    Concatenates all test predictions and important sequences TSV files from the output directory.\n\n    This function finds all files matching the patterns:\n    - *_test_predictions.tsv\n    - *_important_sequences.tsv\n\n    and concatenates them to match the expected output format of submissions.csv.\n\n    Args:\n        out_dir (str): Path to the output directory containing the TSV files.\n\n    Returns:\n        pd.DataFrame: Concatenated DataFrame with predictions followed by important sequences.\n                     Columns: ['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']\n    \"\"\"\n    predictions_pattern = os.path.join(out_dir, '*_test_predictions.tsv')\n    sequences_pattern = os.path.join(out_dir, '*_important_sequences.tsv')\n\n    predictions_files = sorted(glob.glob(predictions_pattern))\n    sequences_files = sorted(glob.glob(sequences_pattern))\n\n    df_list = []\n\n    for pred_file in predictions_files:\n        try:\n            df = pd.read_csv(pred_file, sep='\\t')\n            df_list.append(df)\n        except Exception as e:\n            print(f\"Warning: Could not read predictions file '{pred_file}'. Error: {e}. Skipping.\")\n            continue\n\n    for seq_file in sequences_files:\n        try:\n            df = pd.read_csv(seq_file, sep='\\t')\n            df_list.append(df)\n        except Exception as e:\n            print(f\"Warning: Could not read sequences file '{seq_file}'. Error: {e}. Skipping.\")\n            continue\n\n    if not df_list:\n        print(\"Warning: No output files were found to concatenate.\")\n        concatenated_df = pd.DataFrame(\n            columns=['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call'])\n    else:\n        concatenated_df = pd.concat(df_list, ignore_index=True)\n    submissions_file = os.path.join(out_dir, 'submissions.csv')\n    concatenated_df.to_csv(submissions_file, index=False)\n    print(f\"Concatenated output written to `{submissions_file}`.\")\n\n\ndef get_dataset_pairs(train_dir: str, test_dir: str) -> List[Tuple[str, List[str]]]:\n    \"\"\"Returns list of (train_path, [test_paths]) tuples for dataset pairs.\"\"\"\n    test_groups = defaultdict(list)\n    for test_name in sorted(os.listdir(test_dir)):\n        if test_name.startswith(\"test_dataset_\"):\n            base_id = test_name.replace(\"test_dataset_\", \"\").split(\"_\")[0]\n            test_groups[base_id].append(os.path.join(test_dir, test_name))\n\n    pairs = []\n    for train_name in sorted(os.listdir(train_dir)):\n        if train_name.startswith(\"train_dataset_\"):\n            train_id = train_name.replace(\"train_dataset_\", \"\")\n            train_path = os.path.join(train_dir, train_name)\n            pairs.append((train_path, test_groups.get(train_id, [])))\n\n    return pairs\n","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2025-11-04T13:55:47.790208Z","iopub.execute_input":"2025-11-04T13:55:47.790576Z","iopub.status.idle":"2025-11-04T13:55:47.820888Z","shell.execute_reply.started":"2025-11-04T13:55:47.790558Z","shell.execute_reply":"2025-11-04T13:55:47.819813Z"}},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## A baseline model\n\nBelow, we implement a baseline model that learns and predicts the labels of examples, as well as identifies the important unique rows from the entire training dataset that best explain the labels. The pipeline involves standardizing 3-mer encoded k-mer frequency data, and training a L1-regularized logistic regression model, where the hyperparameter related to regularization strength is chosen through cross-validation. It also implements a method that scores each sequence based on the learnt coefficients to identify the important sequences that are potentially label-associated. We will use the implementation shown in the chunk below in the `ImmuneStatePredictor` class provided in the code template.","metadata":{}},{"cell_type":"code","source":"## A Classifier class that implements functionality of a baseline prediction + identification of sequences that explain the labels\n## This implementation will be used in the provided `ImmuneStatePredictor` class provided through the code template, as shown in the next chunk.\n\n\nclass KmerClassifier:\n    \"\"\"L1-regularized logistic regression for k-mer count data.\"\"\"\n\n    def __init__(self, c_values=None, cv_folds=5,\n                 opt_metric='balanced_accuracy', random_state=123, n_jobs=1):\n        if c_values is None:\n            c_values = [1, 0.1, 0.05, 0.03]\n        self.c_values = c_values\n        self.cv_folds = cv_folds\n        self.opt_metric = opt_metric\n        self.random_state = random_state\n        self.n_jobs = n_jobs\n        self.best_C_ = None\n        self.best_score_ = None\n        self.cv_results_ = None\n        self.model_ = None\n        self.feature_names_ = None\n        self.val_score_ = None\n\n    def _make_pipeline(self, C):\n        \"\"\"Create standardization + L1 logistic regression pipeline.\"\"\"\n        return Pipeline([\n            ('scaler', StandardScaler()),\n            ('classifier', LogisticRegression(\n                penalty='l1', C=C, solver='liblinear',\n                random_state=self.random_state, max_iter=1000\n            ))\n        ])\n\n    def _get_scorer(self):\n        \"\"\"Get scoring function for optimization.\"\"\"\n        if self.opt_metric == 'balanced_accuracy':\n            return 'balanced_accuracy'\n        elif self.opt_metric == 'roc_auc':\n            return 'roc_auc'\n        else:\n            raise ValueError(f\"Unknown metric: {self.opt_metric}\")\n\n    def tune_and_fit(self, X, y, val_size=0.2):\n        \"\"\"Perform CV tuning on train split and fit, with optional validation split.\"\"\"\n\n        if isinstance(X, pd.DataFrame):\n            self.feature_names_ = X.columns.tolist()\n            X = X.values\n        if isinstance(y, pd.Series):\n            y = y.values\n\n        if val_size > 0:\n            X_train, X_val, y_train, y_val = train_test_split(\n                X, y, test_size=val_size, random_state=self.random_state, stratify=y)\n        else:\n            X_train, y_train = X, y\n            X_val, y_val = None, None\n\n        cv = StratifiedKFold(n_splits=self.cv_folds, shuffle=True,\n                             random_state=self.random_state)\n        scorer = self._get_scorer()\n\n        results = []\n        for C in self.c_values:\n            pipeline = self._make_pipeline(C)\n            scores = cross_val_score(pipeline, X_train, y_train, cv=cv, scoring=scorer,\n                                     n_jobs=self.n_jobs)\n            results.append({\n                'C': C,\n                'mean_score': scores.mean(),\n                'std_score': scores.std(),\n                'scores': scores\n            })\n\n        self.cv_results_ = pd.DataFrame(results)\n        best_idx = self.cv_results_['mean_score'].idxmax()\n        self.best_C_ = self.cv_results_.loc[best_idx, 'C']\n        self.best_score_ = self.cv_results_.loc[best_idx, 'mean_score']\n\n        print(f\"Best C: {self.best_C_} (CV {self.opt_metric}: {self.best_score_:.4f})\")\n\n        # Fit on training split with best hyperparameter\n        self.model_ = self._make_pipeline(self.best_C_)\n        self.model_.fit(X_train, y_train)\n\n        if X_val is not None:\n            if scorer == 'balanced_accuracy':\n                self.val_score_ = balanced_accuracy_score(y_val, self.model_.predict(X_val))\n            else:  # roc_auc\n                self.val_score_ = roc_auc_score(y_val, self.model_.predict_proba(X_val)[:, 1])\n            print(f\"Validation {self.opt_metric}: {self.val_score_:.4f}\")\n\n        return self\n\n    def predict_proba(self, X):\n        \"\"\"Predict class probabilities.\"\"\"\n        if self.model_ is None:\n            raise ValueError(\"Model not fitted.\")\n        if isinstance(X, pd.DataFrame):\n            X = X.values\n        return self.model_.predict_proba(X)[:, 1]\n\n    def predict(self, X):\n        \"\"\"Predict class labels.\"\"\"\n        if self.model_ is None:\n            raise ValueError(\"Model not fitted.\")\n        if isinstance(X, pd.DataFrame):\n            X = X.values\n        return self.model_.predict(X)\n\n    def get_feature_importance(self):\n        \"\"\"\n        Get feature importance from L1 coefficients.\n\n        Returns:\n            pd.DataFrame with columns ['feature', 'coefficient', 'abs_coefficient']\n        \"\"\"\n        if self.model_ is None:\n            raise ValueError(\"Model not fitted.\")\n\n        coef = self.model_.named_steps['classifier'].coef_[0]\n\n        if self.feature_names_ is not None:\n            feature_names = self.feature_names_\n        else:\n            feature_names = [f\"feature_{i}\" for i in range(len(coef))]\n\n        importance_df = pd.DataFrame({\n            'feature': feature_names,\n            'coefficient': coef,\n            'abs_coefficient': np.abs(coef)\n        })\n\n        importance_df = importance_df.sort_values('abs_coefficient', ascending=False)\n\n        return importance_df\n\n    def score_all_sequences(self, sequences_df, sequence_col='junction_aa'):\n        \"\"\"\n        Score all sequences using model coefficients.\n\n        Parameters:\n            sequences_df: DataFrame with unique sequences\n            sequence_col: Column name containing sequences\n\n        Returns:\n            DataFrame with added 'importance_score' column\n        \"\"\"\n        if self.model_ is None:\n            raise ValueError(\"Model not fitted.\")\n\n        scaler = self.model_.named_steps['scaler']\n        coefficients = self.model_.named_steps['classifier'].coef_[0]\n        coefficients = coefficients / scaler.scale_\n\n        kmer_to_index = {kmer: idx for idx, kmer in enumerate(self.feature_names_)}\n        k = len(self.feature_names_[0])\n\n        scores = []\n        total_seqs = len(sequences_df)\n        for seq in tqdm(sequences_df[sequence_col], total=total_seqs, desc=\"Scoring sequences\"):\n            counts = np.zeros(len(kmer_to_index), dtype=np.uint8)\n            for i in range(len(seq) - k + 1):\n                kmer = seq[i:i + k]\n                if kmer in kmer_to_index:\n                    counts[kmer_to_index[kmer]] = 1\n            scores.append(np.dot(counts, coefficients))\n\n        result_df = sequences_df.copy()\n        result_df['importance_score'] = scores\n        return result_df\n\n\ndef prepare_data(X_df, labels_df, id_col='ID', label_col='label_positive'):\n    \"\"\"\n    Merge feature matrix with labels, ensuring alignment.\n\n    Parameters:\n        X_df: DataFrame with samples as rows (index contains IDs)\n        labels_df: DataFrame with ID column and label column\n        id_col: Name of ID column in labels_df\n        label_col: Name of label column in labels_df\n\n    Returns:\n        X: Feature matrix aligned with labels\n        y: Binary labels\n        common_ids: IDs that were kept\n    \"\"\"\n    if id_col in labels_df.columns:\n        labels_indexed = labels_df.set_index(id_col)[label_col]\n    else:\n        # Assume labels_df index is already the ID\n        labels_indexed = labels_df[label_col]\n\n    common_ids = X_df.index.intersection(labels_indexed.index)\n\n    if len(common_ids) == 0:\n        raise ValueError(\"No common IDs found between feature matrix and labels\")\n\n    X = X_df.loc[common_ids]\n    y = labels_indexed.loc[common_ids]\n\n    print(f\"Aligned {len(common_ids)} samples with labels\")\n\n    return X, y, common_ids","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## The Predictor Pipeline\n\nThe `ImmuneStatePredictor` class shown below orchestrates the complete workflow adhering to the code template:\n\n1. **Training**: Load examples, encode k-mers, tune hyperparameters via cross-validation, and fit the model\n2. **Prediction**: Generate probability scores for test examples\n3. **Sequence Identification**: Score individual sequences by their k-mer content to find label-associated sequences\n\nThis structure allows the same trained model to both make example-level predictions and identify important sequences. Note that we filled only some lines within the placeholders of the code template.","metadata":{}},{"cell_type":"code","source":"## Main ImmuneStatePredictor class, where this notebook fills in a baseline predictor implementation within the placeholders \n## and replaces any example code lines with actual code that makes sense\n\nclass ImmuneStatePredictor:\n    \"\"\"\n    A template for predicting immune states from TCR repertoire data.\n\n    Participants should implement the logic for training, prediction, and\n    sequence identification within this class.\n    \"\"\"\n\n    def __init__(self, n_jobs: int = 1, device: str = 'cpu', **kwargs):\n        \"\"\"\n        Initializes the predictor.\n\n        Args:\n            n_jobs (int): Number of CPU cores to use for parallel processing.\n            device (str): The device to use for computation (e.g., 'cpu', 'cuda').\n            **kwargs: Additional hyperparameters for the model.\n        \"\"\"\n        self.train_ids_ = None\n        total_cores = os.cpu_count()\n        if n_jobs == -1:\n            self.n_jobs = total_cores\n        else:\n            self.n_jobs = min(n_jobs, total_cores)\n        self.device = device\n        if device == 'cuda' and not torch.cuda.is_available():\n            print(\"Warning: 'cuda' was requested but is not available. Falling back to 'cpu'.\")\n            self.device = 'cpu'\n        else:\n            self.device = device\n        # --- your code starts here ---\n        # Example: Store hyperparameters, the actual model, identified important sequences, etc.\n\n        # NOTE: we encourage you to use self.n_jobs and self.device if appropriate in\n        # your implementation instead of hardcoding these values because your code may later be run in an\n        # environment with different hardware resources.\n\n        self.model = None\n        self.important_sequences_ = None\n        # --- your code ends here ---\n\n    def fit(self, train_dir_path: str):\n        \"\"\"\n        Trains the model on the provided training data.\n\n        Args:\n            train_dir_path (str): Path to the directory with training TSV files.\n\n        Returns:\n            self: The fitted predictor instance.\n        \"\"\"\n\n        # --- your code starts here ---\n        # Load the data, prepare suited representations as needed, train your model,\n        # and find the top k important sequences that best explain the labels.\n        # Example: Load the data. One possibility could be to use the provided utility function as shown below.\n\n        # full_train_dataset_df = load_full_dataset(train_dir_path)\n\n        X_train_df, y_train_df = load_and_encode_kmers(train_dir_path, k=3)  # Example of loading and encoding kmers\n\n        #   Model Training\n        #    Example: self.model = SomeClassifier().fit(X_train, y_train)\n\n        X_train, y_train, train_ids = prepare_data(X_train_df, y_train_df,\n                                                   id_col='ID', label_col='label_positive')\n\n        self.model = KmerClassifier(\n            c_values=[1, 0.2, 0.1, 0.05, 0.03],\n            cv_folds=5,\n            opt_metric='roc_auc',\n            random_state=123,\n            n_jobs=self.n_jobs\n        )\n\n        self.model.tune_and_fit(X_train, y_train)\n\n        self.train_ids_ = train_ids\n\n        #   Identify important sequences (can be done here or in the dedicated method)\n        #    Example:\n        self.important_sequences_ = self.identify_associated_sequences(train_dir_path=train_dir_path, top_k=50000)\n\n        # --- your code ends here ---\n        print(\"Training complete.\")\n        return self\n\n    def predict_proba(self, test_dir_path: str) -> pd.DataFrame:\n        \"\"\"\n        Predicts probabilities for examples in the provided path.\n\n        Args:\n            test_dir_path (str): Path to the directory with test TSV files.\n\n        Returns:\n            pd.DataFrame: A DataFrame with 'ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call' columns.\n        \"\"\"\n        print(f\"Making predictions for data in {test_dir_path}...\")\n        if self.model is None:\n            raise RuntimeError(\"The model has not been fitted yet. Please call `fit` first.\")\n\n        # --- your code starts here ---\n\n        # Example: Load the data. One possibility could be to use the provided utility function as shown below.\n\n        # full_test_dataset_df = load_full_dataset(test_dir_path)\n\n        X_test_df, _ = load_and_encode_kmers(test_dir_path, k=3)\n\n        if self.model.feature_names_ is not None:\n            X_test_df = X_test_df.reindex(columns=self.model.feature_names_, fill_value=0)\n\n        repertoire_ids = X_test_df.index.tolist()\n\n        # Prediction\n        #    Example:\n        # draw random probabilities for demonstration purposes\n\n        probabilities = self.model.predict_proba(X_test_df)\n\n        # --- your code ends here ---\n\n        predictions_df = pd.DataFrame({\n            'ID': repertoire_ids,\n            'dataset': [os.path.basename(test_dir_path)] * len(repertoire_ids),\n            'label_positive_probability': probabilities\n        })\n\n        # to enable compatibility with the expected output format that includes junction_aa, v_call, j_call columns\n        predictions_df['junction_aa'] = -999.0\n        predictions_df['v_call'] = -999.0\n        predictions_df['j_call'] = -999.0\n\n        predictions_df = predictions_df[['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']]\n\n        print(f\"Prediction complete on {len(repertoire_ids)} examples in {test_dir_path}.\")\n        return predictions_df\n\n    def identify_associated_sequences(self, train_dir_path: str, top_k: int = 50000) -> pd.DataFrame:\n        \"\"\"\n        Identifies the top \"k\" important sequences (rows) from the training data that best explain the labels.\n\n        Args:\n            top_k (int): The number of top sequences to return (based on some scoring mechanism).\n            train_dir_path (str): Path to the directory with training TSV files.\n\n        Returns:\n            pd.DataFrame: A DataFrame with 'ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call' columns.\n        \"\"\"\n        dataset_name = os.path.basename(train_dir_path)\n\n        # --- your code starts here ---\n        # Return the top k sequences, sorted based on some form of importance score.\n        # Example:\n        # all_sequences_scored = self._score_all_sequences()\n\n        full_df = load_full_dataset(train_dir_path)\n        unique_seqs = full_df[['junction_aa', 'v_call', 'j_call']].drop_duplicates()\n        all_sequences_scored = self.model.score_all_sequences(unique_seqs, sequence_col='junction_aa')\n\n        # --- your code ends here ---\n\n        top_sequences_df = all_sequences_scored.nlargest(top_k, 'importance_score')\n        top_sequences_df = top_sequences_df[['junction_aa', 'v_call', 'j_call']]\n        top_sequences_df['dataset'] = dataset_name\n        top_sequences_df['ID'] = range(1, len(top_sequences_df)+1)\n        top_sequences_df['ID'] = top_sequences_df['dataset'] + '_seq_top_' + top_sequences_df['ID'].astype(str)\n        top_sequences_df['label_positive_probability'] = -999.0 # to enable compatibility with the expected output format\n        top_sequences_df = top_sequences_df[['ID', 'dataset', 'label_positive_probability', 'junction_aa', 'v_call', 'j_call']]\n\n        return top_sequences_df\n","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Execution Workflow\n\nThe functions below from the code template handle the end-to-end workflow for one training + test data combination: training on each dataset, generating predictions for test sets, and saving results in the required submission format. We are not changing anything from the template.","metadata":{}},{"cell_type":"code","source":"## The `main` workflow that uses your implementation of the ImmuneStatePredictor class to train, identify important sequences and predict test labels\n\n\ndef _train_predictor(predictor: ImmuneStatePredictor, train_dir: str):\n    \"\"\"Trains the predictor on the training data.\"\"\"\n    print(f\"Fitting model on examples in ` {train_dir} `...\")\n    predictor.fit(train_dir)\n\n\ndef _generate_predictions(predictor: ImmuneStatePredictor, test_dirs: List[str]) -> pd.DataFrame:\n    \"\"\"Generates predictions for all test directories and concatenates them.\"\"\"\n    all_preds = []\n    for test_dir in test_dirs:\n        print(f\"Predicting on examples in ` {test_dir} `...\")\n        preds = predictor.predict_proba(test_dir)\n        if preds is not None and not preds.empty:\n            all_preds.append(preds)\n        else:\n            print(f\"Warning: No predictions returned for {test_dir}\")\n    if all_preds:\n        return pd.concat(all_preds, ignore_index=True)\n    return pd.DataFrame()\n\n\ndef _save_predictions(predictions: pd.DataFrame, out_dir: str, train_dir: str) -> None:\n    \"\"\"Saves predictions to a TSV file.\"\"\"\n    if predictions.empty:\n        raise ValueError(\"No predictions to save - predictions DataFrame is empty\")\n\n    preds_path = os.path.join(out_dir, f\"{os.path.basename(train_dir)}_test_predictions.tsv\")\n    save_tsv(predictions, preds_path)\n    print(f\"Predictions written to `{preds_path}`.\")\n\n\ndef _save_important_sequences(predictor: ImmuneStatePredictor, out_dir: str, train_dir: str) -> None:\n    \"\"\"Saves important sequences to a TSV file.\"\"\"\n    seqs = predictor.important_sequences_\n    if seqs is None or seqs.empty:\n        raise ValueError(\"No important sequences available to save\")\n\n    seqs_path = os.path.join(out_dir, f\"{os.path.basename(train_dir)}_important_sequences.tsv\")\n    save_tsv(seqs, seqs_path)\n    print(f\"Important sequences written to `{seqs_path}`.\")\n\n\ndef main(train_dir: str, test_dirs: List[str], out_dir: str, n_jobs: int, device: str) -> None:\n    validate_dirs_and_files(train_dir, test_dirs, out_dir)\n    predictor = ImmuneStatePredictor(n_jobs=n_jobs,\n                                     device=device)  # instantiate with any other parameters as defined by you in the class\n    _train_predictor(predictor, train_dir)\n    predictions = _generate_predictions(predictor, test_dirs)\n    _save_predictions(predictions, out_dir, train_dir)\n    _save_important_sequences(predictor, out_dir, train_dir)\n\n\ndef run():\n    parser = argparse.ArgumentParser(description=\"Immune State Predictor CLI\")\n    parser.add_argument(\"--train_dir\", required=True, help=\"Path to training data directory\")\n    parser.add_argument(\"--test_dirs\", required=True, nargs=\"+\", help=\"Path(s) to test data director(ies)\")\n    parser.add_argument(\"--out_dir\", required=True, help=\"Path to output directory\")\n    parser.add_argument(\"--n_jobs\", type=int, default=1,\n                        help=\"Number of CPU cores to use. Use -1 for all available cores.\")\n    parser.add_argument(\"--device\", type=str, default='cpu', choices=['cpu', 'cuda'],\n                        help=\"Device to use for computation ('cpu' or 'cuda').\")\n    args = parser.parse_args()\n    main(args.train_dir, args.test_dirs, args.out_dir, args.n_jobs, args.device)","metadata":{"trusted":true},"outputs":[],"execution_count":null},{"cell_type":"markdown","source":"## Running the Pipeline\n\nNow we execute the complete workflow on all dataset pairs. Each training dataset has corresponding test sets that our model will predict on and saves results. Finally, we will use a utility function `concatenate_output_files` that writes a `submission.csv` to disk.","metadata":{}},{"cell_type":"code","source":"train_datasets_dir = \"/kaggle/input/adaptive-immune-profiling-challenge-2025/train_datasets/train_datasets\"\ntest_datasets_dir = \"/kaggle/input/adaptive-immune-profiling-challenge-2025/test_datasets/test_datasets\"\nresults_dir = \"/kaggle/working/results\"\n\ntrain_test_dataset_pairs = get_dataset_pairs(train_datasets_dir, test_datasets_dir)\n\nfor train_dir, test_dirs in train_test_dataset_pairs:\n    main(train_dir=train_dir, test_dirs=test_dirs, out_dir=results_dir, n_jobs=4, device=\"cpu\")\n\nconcatenate_output_files(results_dir)","metadata":{"trusted":true},"outputs":[],"execution_count":null}]}