TCRmeta 0.1.0__tar.gz

tcrmeta-0.1.0/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: TCRmeta
+Version: 0.1.0
+Summary: Antigen-aware TCR repertoire embeddings: pretrained ESM2 + contrastive fine-tuning, with downstream CSS scoring, UMAP projection, and energy-distance shift analysis.
+Author: Mingyao Pan
+License: MIT
+Project-URL: Homepage, https://github.com/Mia-yao/TCRmeta
+Keywords: TCR,immunology,repertoire,embedding,ESM2,contrastive learning
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Requires-Dist: scipy>=1.9
+Requires-Dist: scikit-learn>=1.2
+Requires-Dist: torch>=1.13
+Requires-Dist: transformers>=4.30
+Requires-Dist: umap-learn>=0.5
+Requires-Dist: joblib>=1.2
+Requires-Dist: tqdm>=4.64
+Requires-Dist: huggingface_hub>=0.20
+Requires-Dist: matplotlib>=3.6
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: build>=1.0; extra == "dev"
+
+# TCRmeta
+
+Antigen-aware TCR repertoire embeddings and downstream repertoire-level
+analysis, built on a pretrained-ESM2 + contrastively fine-tuned encoder.
+
+Given a bulk TCR repertoire, TCRmeta can:
+
+1. **Embed** every clone into a TCR embedding (`embed_repertoire`) — pick between a "pretrained" or a "final" embedding, see below.
+2. **Score** a repertoire's clonal shift score (CSS) against a reference cohort (`compute_css`).
+3. **Project** a repertoire onto a reference UMAP map, per V gene (`plot_umap`).
+4. **Compare** two repertoires via a frequency-weighted energy-distance shift (`energy_shift`).
+
+Model weights and the default reference map are downloaded once from the
+Hugging Face Hub and cached locally (`~/.cache/tcrmeta` by default).
+
+## Install
+
+```bash
+pip install tcrmeta
+```
+
+GPU acceleration (embedding + energy-distance computation) is used
+automatically if a CUDA (or MPS) device is available; pass `device="cpu"`
+anywhere to force CPU.
+
+## Input format
+
+Every function takes a TCR repertoire as a `pandas.DataFrame` with
+**exactly** these columns (no aliases are auto-detected — rename your own
+columns first):
+
+| column   | meaning                                  |
+|----------|-------------------------------------------|
+| `cdr3aa` | CDR3 amino acid sequence                  |
+| `v_gene` | V gene (e.g. `TRBV6-4`, allele suffix ok)  |
+| `count`  | clone count / templates / reads           |
+
+Currently supports human TRB (beta chain) repertoires only.
+
+## Quickstart
+
+```python
+import pandas as pd
+import tcrmeta as tm
+
+df = pd.read_csv("my_repertoire.csv")  # cdr3aa, v_gene, count
+
+# 1. Embed (embedding_type="final" is the default — see "Choosing an
+# embedding type" below for when to use "pretrained" instead)
+embeddings = tm.embed_repertoire(df)  # dict {(cdr3aa, v_gene): 64-dim np.ndarray}
+
+# 2. CSS against the shipped default reference
+per_gene_css, whole_css = tm.compute_css(df)  # v_gene=None -> scores all reference V genes
+
+# 3. UMAP projection (v_gene required)
+fig, coords = tm.plot_umap(df, v_gene="TRBV6-4")
+fig.savefig("trbv6-4_umap.png")
+
+# 4. Energy-distance shift between two repertoires
+df2 = pd.read_csv("other_repertoire.csv")
+per_gene_energy, weighted_summary = tm.energy_shift(df, df2)
+```
+
+### Building your own reference map
+
+```python
+reference_repertoires = [pd.read_csv(f) for f in my_reference_files]
+reference = tm.build_reference(reference_repertoires)
+tm.save_reference(reference, "my_reference.pkl")
+
+per_gene_css, whole_css = tm.compute_css(df, reference="my_reference.pkl")
+fig, coords = tm.plot_umap(df, reference=reference, v_gene="TRBV6-4")
+```
+
+### Choosing an embedding type
+
+`embed_repertoire` accepts an `embedding_type` argument with two options:
+
+| `embedding_type` | Dim | What it is | What it captures |
+|---|---|---|---|
+| `"pretrained"` | 480 | The raw CLS embedding straight out of the masked-language-model-pretrained ESM2-style base encoder, before any contrastive fine-tuning. | **Local structure** — this encoder is trained to recover masked residues from local sequence context, so the embedding is most sensitive to motif/sub-sequence-level similarity between TCRs. |
+| `"final"` (default) | 64 | The 480-dim base embedding run through the ensemble of 7 contrastively fine-tuned projection heads, GPA-aligned and mean-fused, L2-normalized. | **Overall structure** — contrastive fine-tuning pulls together TCRs recognizing the same antigen regardless of local sequence differences, so this embedding is most sensitive to antigen-specificity-level, global similarity. |
+
+```python
+# Local-structure embedding
+emb_pretrained = tm.embed_repertoire(df, embedding_type="pretrained")
+
+# Overall-structure embedding (default; same as tm.embed_repertoire(df))
+emb_final = tm.embed_repertoire(df, embedding_type="final")
+```
+
+This choice is only exposed on `embed_repertoire` itself, for users who
+want the raw embeddings for their own downstream analysis.
+`compute_css`, `plot_umap`, `energy_shift`, and `build_reference` always
+embed internally with `embedding_type="final"` — TCRmeta's own
+repertoire-level statistics (and the shipped reference map) are all
+defined against that 64-dim antigen-aware embedding space, so these
+functions ignore any `embedding_type` passed to them.
+
+### Keeping intermediate embeddings
+
+```python
+tm.embed_repertoire(df, save_path="my_repertoire_embeddings.pkl")
+```
+
+Pass the saved path back in via `embeddings=` (loaded with `pickle.load`)
+to any downstream function to skip re-embedding, or just delete the file
+if you don't need it.
+
+## Configuration
+
+| Env var               | Purpose                                              |
+|------------------------|-------------------------------------------------------|
+| `TCRMETA_HF_REPO`      | Hugging Face Hub repo id for weights/reference (placeholder until published) |
+| `TCRMETA_WEIGHTS_DIR`  | Load weights from a local directory instead of the Hub |
+| `TCRMETA_CACHE_DIR`    | Local cache directory (default `~/.cache/tcrmeta`)    |
+
+See `scripts/upload_weights.py` for a one-time script to publish your
+model weights and default reference map to the Hugging Face Hub.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest
+```

tcrmeta-0.1.0/README.md

@@ -0,0 +1,126 @@
+# TCRmeta
+
+Antigen-aware TCR repertoire embeddings and downstream repertoire-level
+analysis, built on a pretrained-ESM2 + contrastively fine-tuned encoder.
+
+Given a bulk TCR repertoire, TCRmeta can:
+
+1. **Embed** every clone into a TCR embedding (`embed_repertoire`) — pick between a "pretrained" or a "final" embedding, see below.
+2. **Score** a repertoire's clonal shift score (CSS) against a reference cohort (`compute_css`).
+3. **Project** a repertoire onto a reference UMAP map, per V gene (`plot_umap`).
+4. **Compare** two repertoires via a frequency-weighted energy-distance shift (`energy_shift`).
+
+Model weights and the default reference map are downloaded once from the
+Hugging Face Hub and cached locally (`~/.cache/tcrmeta` by default).
+
+## Install
+
+```bash
+pip install tcrmeta
+```
+
+GPU acceleration (embedding + energy-distance computation) is used
+automatically if a CUDA (or MPS) device is available; pass `device="cpu"`
+anywhere to force CPU.
+
+## Input format
+
+Every function takes a TCR repertoire as a `pandas.DataFrame` with
+**exactly** these columns (no aliases are auto-detected — rename your own
+columns first):
+
+| column   | meaning                                  |
+|----------|-------------------------------------------|
+| `cdr3aa` | CDR3 amino acid sequence                  |
+| `v_gene` | V gene (e.g. `TRBV6-4`, allele suffix ok)  |
+| `count`  | clone count / templates / reads           |
+
+Currently supports human TRB (beta chain) repertoires only.
+
+## Quickstart
+
+```python
+import pandas as pd
+import tcrmeta as tm
+
+df = pd.read_csv("my_repertoire.csv")  # cdr3aa, v_gene, count
+
+# 1. Embed (embedding_type="final" is the default — see "Choosing an
+# embedding type" below for when to use "pretrained" instead)
+embeddings = tm.embed_repertoire(df)  # dict {(cdr3aa, v_gene): 64-dim np.ndarray}
+
+# 2. CSS against the shipped default reference
+per_gene_css, whole_css = tm.compute_css(df)  # v_gene=None -> scores all reference V genes
+
+# 3. UMAP projection (v_gene required)
+fig, coords = tm.plot_umap(df, v_gene="TRBV6-4")
+fig.savefig("trbv6-4_umap.png")
+
+# 4. Energy-distance shift between two repertoires
+df2 = pd.read_csv("other_repertoire.csv")
+per_gene_energy, weighted_summary = tm.energy_shift(df, df2)
+```
+
+### Building your own reference map
+
+```python
+reference_repertoires = [pd.read_csv(f) for f in my_reference_files]
+reference = tm.build_reference(reference_repertoires)
+tm.save_reference(reference, "my_reference.pkl")
+
+per_gene_css, whole_css = tm.compute_css(df, reference="my_reference.pkl")
+fig, coords = tm.plot_umap(df, reference=reference, v_gene="TRBV6-4")
+```
+
+### Choosing an embedding type
+
+`embed_repertoire` accepts an `embedding_type` argument with two options:
+
+| `embedding_type` | Dim | What it is | What it captures |
+|---|---|---|---|
+| `"pretrained"` | 480 | The raw CLS embedding straight out of the masked-language-model-pretrained ESM2-style base encoder, before any contrastive fine-tuning. | **Local structure** — this encoder is trained to recover masked residues from local sequence context, so the embedding is most sensitive to motif/sub-sequence-level similarity between TCRs. |
+| `"final"` (default) | 64 | The 480-dim base embedding run through the ensemble of 7 contrastively fine-tuned projection heads, GPA-aligned and mean-fused, L2-normalized. | **Overall structure** — contrastive fine-tuning pulls together TCRs recognizing the same antigen regardless of local sequence differences, so this embedding is most sensitive to antigen-specificity-level, global similarity. |
+
+```python
+# Local-structure embedding
+emb_pretrained = tm.embed_repertoire(df, embedding_type="pretrained")
+
+# Overall-structure embedding (default; same as tm.embed_repertoire(df))
+emb_final = tm.embed_repertoire(df, embedding_type="final")
+```
+
+This choice is only exposed on `embed_repertoire` itself, for users who
+want the raw embeddings for their own downstream analysis.
+`compute_css`, `plot_umap`, `energy_shift`, and `build_reference` always
+embed internally with `embedding_type="final"` — TCRmeta's own
+repertoire-level statistics (and the shipped reference map) are all
+defined against that 64-dim antigen-aware embedding space, so these
+functions ignore any `embedding_type` passed to them.
+
+### Keeping intermediate embeddings
+
+```python
+tm.embed_repertoire(df, save_path="my_repertoire_embeddings.pkl")
+```
+
+Pass the saved path back in via `embeddings=` (loaded with `pickle.load`)
+to any downstream function to skip re-embedding, or just delete the file
+if you don't need it.
+
+## Configuration
+
+| Env var               | Purpose                                              |
+|------------------------|-------------------------------------------------------|
+| `TCRMETA_HF_REPO`      | Hugging Face Hub repo id for weights/reference (placeholder until published) |
+| `TCRMETA_WEIGHTS_DIR`  | Load weights from a local directory instead of the Hub |
+| `TCRMETA_CACHE_DIR`    | Local cache directory (default `~/.cache/tcrmeta`)    |
+
+See `scripts/upload_weights.py` for a one-time script to publish your
+model weights and default reference map to the Hugging Face Hub.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest
+```

tcrmeta-0.1.0/pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "TCRmeta"
+version = "0.1.0"
+description = "Antigen-aware TCR repertoire embeddings: pretrained ESM2 + contrastive fine-tuning, with downstream CSS scoring, UMAP projection, and energy-distance shift analysis."
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "MIT" }
+authors = [
+    { name = "Mingyao Pan" },
+]
+keywords = ["TCR", "immunology", "repertoire", "embedding", "ESM2", "contrastive learning"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Intended Audience :: Science/Research",
+    "Topic :: Scientific/Engineering :: Bio-Informatics",
+]
+
+dependencies = [
+    "numpy>=1.23",
+    "pandas>=1.5",
+    "scipy>=1.9",
+    "scikit-learn>=1.2",
+    "torch>=1.13",
+    "transformers>=4.30",
+    "umap-learn>=0.5",
+    "joblib>=1.2",
+    "tqdm>=4.64",
+    "huggingface_hub>=0.20",
+    "matplotlib>=3.6",
+]
+
+[project.optional-dependencies]
+dev = ["pytest>=7.0", "build>=1.0"]
+
+[project.urls]
+Homepage = "https://github.com/Mia-yao/TCRmeta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.setuptools.package-data]
+tcrmeta = ["data/*.csv"]

tcrmeta-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

tcrmeta-0.1.0/src/TCRmeta.egg-info/PKG-INFO

@@ -0,0 +1,154 @@
+Metadata-Version: 2.4
+Name: TCRmeta
+Version: 0.1.0
+Summary: Antigen-aware TCR repertoire embeddings: pretrained ESM2 + contrastive fine-tuning, with downstream CSS scoring, UMAP projection, and energy-distance shift analysis.
+Author: Mingyao Pan
+License: MIT
+Project-URL: Homepage, https://github.com/Mia-yao/TCRmeta
+Keywords: TCR,immunology,repertoire,embedding,ESM2,contrastive learning
+Classifier: Programming Language :: Python :: 3
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.23
+Requires-Dist: pandas>=1.5
+Requires-Dist: scipy>=1.9
+Requires-Dist: scikit-learn>=1.2
+Requires-Dist: torch>=1.13
+Requires-Dist: transformers>=4.30
+Requires-Dist: umap-learn>=0.5
+Requires-Dist: joblib>=1.2
+Requires-Dist: tqdm>=4.64
+Requires-Dist: huggingface_hub>=0.20
+Requires-Dist: matplotlib>=3.6
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+Requires-Dist: build>=1.0; extra == "dev"
+
+# TCRmeta
+
+Antigen-aware TCR repertoire embeddings and downstream repertoire-level
+analysis, built on a pretrained-ESM2 + contrastively fine-tuned encoder.
+
+Given a bulk TCR repertoire, TCRmeta can:
+
+1. **Embed** every clone into a TCR embedding (`embed_repertoire`) — pick between a "pretrained" or a "final" embedding, see below.
+2. **Score** a repertoire's clonal shift score (CSS) against a reference cohort (`compute_css`).
+3. **Project** a repertoire onto a reference UMAP map, per V gene (`plot_umap`).
+4. **Compare** two repertoires via a frequency-weighted energy-distance shift (`energy_shift`).
+
+Model weights and the default reference map are downloaded once from the
+Hugging Face Hub and cached locally (`~/.cache/tcrmeta` by default).
+
+## Install
+
+```bash
+pip install tcrmeta
+```
+
+GPU acceleration (embedding + energy-distance computation) is used
+automatically if a CUDA (or MPS) device is available; pass `device="cpu"`
+anywhere to force CPU.
+
+## Input format
+
+Every function takes a TCR repertoire as a `pandas.DataFrame` with
+**exactly** these columns (no aliases are auto-detected — rename your own
+columns first):
+
+| column   | meaning                                  |
+|----------|-------------------------------------------|
+| `cdr3aa` | CDR3 amino acid sequence                  |
+| `v_gene` | V gene (e.g. `TRBV6-4`, allele suffix ok)  |
+| `count`  | clone count / templates / reads           |
+
+Currently supports human TRB (beta chain) repertoires only.
+
+## Quickstart
+
+```python
+import pandas as pd
+import tcrmeta as tm
+
+df = pd.read_csv("my_repertoire.csv")  # cdr3aa, v_gene, count
+
+# 1. Embed (embedding_type="final" is the default — see "Choosing an
+# embedding type" below for when to use "pretrained" instead)
+embeddings = tm.embed_repertoire(df)  # dict {(cdr3aa, v_gene): 64-dim np.ndarray}
+
+# 2. CSS against the shipped default reference
+per_gene_css, whole_css = tm.compute_css(df)  # v_gene=None -> scores all reference V genes
+
+# 3. UMAP projection (v_gene required)
+fig, coords = tm.plot_umap(df, v_gene="TRBV6-4")
+fig.savefig("trbv6-4_umap.png")
+
+# 4. Energy-distance shift between two repertoires
+df2 = pd.read_csv("other_repertoire.csv")
+per_gene_energy, weighted_summary = tm.energy_shift(df, df2)
+```
+
+### Building your own reference map
+
+```python
+reference_repertoires = [pd.read_csv(f) for f in my_reference_files]
+reference = tm.build_reference(reference_repertoires)
+tm.save_reference(reference, "my_reference.pkl")
+
+per_gene_css, whole_css = tm.compute_css(df, reference="my_reference.pkl")
+fig, coords = tm.plot_umap(df, reference=reference, v_gene="TRBV6-4")
+```
+
+### Choosing an embedding type
+
+`embed_repertoire` accepts an `embedding_type` argument with two options:
+
+| `embedding_type` | Dim | What it is | What it captures |
+|---|---|---|---|
+| `"pretrained"` | 480 | The raw CLS embedding straight out of the masked-language-model-pretrained ESM2-style base encoder, before any contrastive fine-tuning. | **Local structure** — this encoder is trained to recover masked residues from local sequence context, so the embedding is most sensitive to motif/sub-sequence-level similarity between TCRs. |
+| `"final"` (default) | 64 | The 480-dim base embedding run through the ensemble of 7 contrastively fine-tuned projection heads, GPA-aligned and mean-fused, L2-normalized. | **Overall structure** — contrastive fine-tuning pulls together TCRs recognizing the same antigen regardless of local sequence differences, so this embedding is most sensitive to antigen-specificity-level, global similarity. |
+
+```python
+# Local-structure embedding
+emb_pretrained = tm.embed_repertoire(df, embedding_type="pretrained")
+
+# Overall-structure embedding (default; same as tm.embed_repertoire(df))
+emb_final = tm.embed_repertoire(df, embedding_type="final")
+```
+
+This choice is only exposed on `embed_repertoire` itself, for users who
+want the raw embeddings for their own downstream analysis.
+`compute_css`, `plot_umap`, `energy_shift`, and `build_reference` always
+embed internally with `embedding_type="final"` — TCRmeta's own
+repertoire-level statistics (and the shipped reference map) are all
+defined against that 64-dim antigen-aware embedding space, so these
+functions ignore any `embedding_type` passed to them.
+
+### Keeping intermediate embeddings
+
+```python
+tm.embed_repertoire(df, save_path="my_repertoire_embeddings.pkl")
+```
+
+Pass the saved path back in via `embeddings=` (loaded with `pickle.load`)
+to any downstream function to skip re-embedding, or just delete the file
+if you don't need it.
+
+## Configuration
+
+| Env var               | Purpose                                              |
+|------------------------|-------------------------------------------------------|
+| `TCRMETA_HF_REPO`      | Hugging Face Hub repo id for weights/reference (placeholder until published) |
+| `TCRMETA_WEIGHTS_DIR`  | Load weights from a local directory instead of the Hub |
+| `TCRMETA_CACHE_DIR`    | Local cache directory (default `~/.cache/tcrmeta`)    |
+
+See `scripts/upload_weights.py` for a one-time script to publish your
+model weights and default reference map to the Hugging Face Hub.
+
+## Development
+
+```bash
+pip install -e ".[dev]"
+pytest
+```

tcrmeta-0.1.0/src/TCRmeta.egg-info/SOURCES.txt

@@ -0,0 +1,22 @@
+README.md
+pyproject.toml
+src/TCRmeta.egg-info/PKG-INFO
+src/TCRmeta.egg-info/SOURCES.txt
+src/TCRmeta.egg-info/dependency_links.txt
+src/TCRmeta.egg-info/requires.txt
+src/TCRmeta.egg-info/top_level.txt
+src/tcrmeta/__init__.py
+src/tcrmeta/_dataset.py
+src/tcrmeta/_model.py
+src/tcrmeta/_tokenizer.py
+src/tcrmeta/_utils.py
+src/tcrmeta/_weights.py
+src/tcrmeta/css.py
+src/tcrmeta/embedding.py
+src/tcrmeta/energy.py
+src/tcrmeta/reference.py
+src/tcrmeta/umap_plot.py
+tests/test_energy.py
+tests/test_reference_and_css.py
+tests/test_umap_plot.py
+tests/test_utils.py

tcrmeta-0.1.0/src/TCRmeta.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

tcrmeta-0.1.0/src/TCRmeta.egg-info/requires.txt

@@ -0,0 +1,15 @@
+numpy>=1.23
+pandas>=1.5
+scipy>=1.9
+scikit-learn>=1.2
+torch>=1.13
+transformers>=4.30
+umap-learn>=0.5
+joblib>=1.2
+tqdm>=4.64
+huggingface_hub>=0.20
+matplotlib>=3.6
+
+[dev]
+pytest>=7.0
+build>=1.0

tcrmeta-0.1.0/src/TCRmeta.egg-info/top_level.txt

@@ -0,0 +1 @@
+tcrmeta

tcrmeta-0.1.0/src/tcrmeta/__init__.py

@@ -0,0 +1,22 @@
+"""TCRmeta: antigen-aware TCR repertoire embeddings and downstream
+repertoire-level analysis (CSS scoring, UMAP projection, energy-distance
+shift), built on a pretrained-ESM2 + contrastive-fine-tuned encoder.
+"""
+from .embedding import embed_repertoire
+from .reference import build_reference, load_reference, save_reference
+from .css import compute_css, load_default_reference
+from .umap_plot import plot_umap
+from .energy import energy_shift
+
+__all__ = [
+    "embed_repertoire",
+    "build_reference",
+    "save_reference",
+    "load_reference",
+    "load_default_reference",
+    "compute_css",
+    "plot_umap",
+    "energy_shift",
+]
+
+__version__ = "0.1.0"

tcrmeta-0.1.0/src/tcrmeta/_dataset.py

@@ -0,0 +1,53 @@
+"""Fixed-length tokenization of (CDR1 + CDR2 + CDR2.5 + CDR3) TCR
+sequences for the ESM2 base encoder.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Tuple
+
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+
+DEFAULT_MAX_LEN = 48
+
+
+class TCRDatasetFixedLen(Dataset):
+    def __init__(self, df: pd.DataFrame, token2idx: Dict[str, int], max_len: int = DEFAULT_MAX_LEN):
+        self.df = df.reset_index(drop=True)
+        self.t2i = token2idx
+        self.max_len = max_len
+        self.cls = token2idx["[CLS]"]
+        self.pad = token2idx["[PAD]"]
+        self.eos = token2idx["[EOS]"]
+        self.unk = token2idx["[UNK]"]
+        self.gap = token2idx["X"]  # separator between CDR segments
+
+    def _encode_seq(self, s: str) -> List[int]:
+        return [self.t2i.get(a, self.unk) for a in s]
+
+    def __getitem__(self, i) -> Tuple[Tuple[str, str], List[int]]:
+        r = self.df.iloc[i]
+        v_seq = r.cdr1 + "X" + r.cdr2 + "X" + r.cdr2_5
+        tokens = (
+            [self.cls]
+            + self._encode_seq(v_seq)
+            + [self.gap]
+            + self._encode_seq(r.cdr3aa)
+            + [self.eos]
+        )
+        if len(tokens) < self.max_len:
+            tokens = tokens + [self.pad] * (self.max_len - len(tokens))
+        else:
+            tokens = tokens[: self.max_len]
+        key = (r.cdr3aa, r.v_gene)
+        return key, tokens
+
+    def __len__(self) -> int:
+        return len(self.df)
+
+
+def fixedlen_collate(batch):
+    keys, seqs = zip(*batch)
+    input_ids = torch.tensor(seqs, dtype=torch.long)
+    return keys, input_ids