hf-tcr

HuggingFace-based inference and evaluation library for TCR-pMHC sequence translation models.

Installation

pip install hf-tcr

Or install from source:

git clone https://github.com/pirl-unc/hf-tcr.git
cd hf-tcr
pip install .

Quick Start

Loading Data

from hf_tcr import TCRpMHCdataset

# Create dataset for pMHC -> TCR translation
dataset = TCRpMHCdataset(
    source="pmhc",
    target="tcr",
    use_pseudo=True,
    use_cdr3=True
)

# Load from CSV file
dataset.load_data_from_file("path/to/data.csv")

Running Inference

from hf_tcr import HuggingFaceModelAdapter, TCRBartTokenizer
from transformers import BartForConditionalGeneration

# Load your trained model and tokenizer
tokenizer = TCRBartTokenizer()
model = BartForConditionalGeneration.from_pretrained("path/to/model")

# Create adapter
adapter = HuggingFaceModelAdapter(
    hf_tokenizer=tokenizer,
    hf_model=model,
    device="cuda"
)

# Get a source from your dataset
source = dataset[0][0]  # Get source from first example

# Generate translations
translations = adapter.sample_translations(
    source=source,
    n=10,
    max_len=25,
    mode="top_k",
    top_k=50,
    temperature=1.0
)

Evaluating Models

from hf_tcr import ModelEvaluator

# Create evaluator (extends HuggingFaceModelAdapter)
evaluator = ModelEvaluator(
    hf_tokenizer=tokenizer,
    hf_model=model,
    device="cuda"
)

# Compute dataset-level metrics
metrics = evaluator.dataset_metrics_at_k(
    dataset=dataset,
    k=100,
    max_len=25,
    mode="top_k",
    top_k=50
)

print(f"BLEU: {metrics['char-bleu']:.4f}")
print(f"Precision@100: {metrics['precision']:.4f}")
print(f"Recall@100: {metrics['recall']:.4f}")
print(f"F1@100: {metrics['f1']:.4f}")
print(f"Mean Edit Distance: {metrics['d_edit']:.2f}")
print(f"Sequence Recovery: {metrics['seq_recovery']:.4f}")
print(f"Diversity: {metrics['diversity']:.4f}")
print(f"Perplexity: {metrics['perplexity']:.2f}")

Available Decoding Strategies

The adapter supports multiple decoding strategies:

• greedy: Deterministic greedy decoding
• ancestral: Multinomial sampling
• top_k: Top-k sampling with temperature
• top_p: Nucleus (top-p) sampling
• beam: Deterministic beam search
• stochastic_beam: Stochastic beam search
• diverse_beam: Diverse beam search
• contrastive: Contrastive decoding
• typical: Typical sampling

Metrics

The ModelEvaluator provides the following metrics:

• Char-BLEU: Character-level BLEU score
• Precision@K: Fraction of generated sequences that match references
• Recall@K: Fraction of reference sequences recovered
• F1@K: Harmonic mean of precision and recall
• Mean Edit Distance: Average Levenshtein distance to closest reference
• Sequence Recovery: Position-wise match percentage
• Diversity: Ratio of unique to total generated sequences
• Perplexity: Model perplexity on the dataset

Soft matching and reference coverage

from hf_tcr import ModelEvaluator

# No model is needed to score existing outputs.
scores = ModelEvaluator.soft_metrics_at_k(
    ["AAAAAAAAAA", "AAAAAAAAAB"], ["AAAAAAAAAA", "CCCCCCCCCC"], k=2,
)
print(scores["soft_precision"], scores["soft_recall"],
      scores["soft_f1"], scores["reference_coverage"])
# 1.0 0.5 0.6666666666666666 0.5

# With an initialized evaluator and dataset, the same keys are macro-averaged:
metrics = evaluator.dataset_metrics_at_k(dataset, k=100, soft_threshold=0.90)
print(metrics["soft_f1"], metrics["reference_coverage"])

Similarity is 1 - Levenshtein(a,b) / max(len(a),len(b)). Similarity is 0 whenever either sequence is empty, including empty/empty, preserving the internal convention. Empty strings cannot match at the default threshold. Recall uses edges with similarity >= 0.90 by default and finds a maximum-cardinality one-to-one matching after thresholding. It does not optimize summed similarity.

For the first k outputs, let G be their count, R the reference count, H the number of outputs hitting any reference, M the number of one-to-one pairs, and C the number of references hit by any evaluated output:

Zero denominators and empty inputs yield zero. All four metrics are bounded by 1. All four scores share one similarity matrix per source; dataset F1 is the mean of per-source F1 scores.

Direct calls do not deduplicate: repeated strings are separate matching vertices and coverage counts reference entries. Dataset conversion retains its existing per-source reference deduplication, while generated entries remain unchanged. Underfilled outputs use actual G for precision and requested k for recall. k=None in direct calls evaluates all outputs and uses their count as k; k=0 returns zeros, and negative k raises ValueError. The dataset API requires a positive integer k because it also generates outputs (None raises TypeError; zero/negative values raise ValueError). Empty datasets return zeros.

Existing exact-match precision, recall, and f1 keys and defaults are unchanged, including their historical behavior on overfilled output lists. As of 0.2.3, soft_precision measures whether each output is near a known reference. Each output counts at most once, but several outputs may match the same reference. soft_recall measures distinct reference recovery using maximum-cardinality one-to-one matching; soft_f1 is their harmonic mean. For identical inputs, precision and F1 may increase relative to 0.2.2; recall and reference coverage are unchanged. Previously, many-to-many soft recall counted all hit references but divided by min(k, R), so a single output could push recall and F1 above 1. One-to-one matching fixes this by limiting each output and reference to one pair. Historical and new soft values are not directly interchangeable. reference_coverage retains many-to-many reference credit with an uncapped reference denominator. No legacy implementations or keys are provided. The internal helper signatures are preserved as thin wrappers around the corrected shared implementation (no historical matching logic is retained):

ModelEvaluator._soft_precision_at_k(outputs, references, threshold=0.90, k=None)
ModelEvaluator._soft_recall_at_k(outputs, references, threshold=0.90, k=100)
ModelEvaluator._soft_f1_at_k(outputs, references, threshold=0.90, k=100)

These helpers preserve the positional argument order threshold, k, while the public methods and soft_metrics_at_k use k, threshold. Prefer explicit keywords when switching between these APIs. Dataset metric keys remain unchanged, with reference_coverage added.

Data Format

CSV files should contain the following columns:

Required:

• CDR3b: CDR3 beta sequence
• TRBV: TRBV gene (IMGT format)
• TRBJ: TRBJ gene (IMGT format)
• Epitope: Peptide sequence
• Allele: HLA allele
• Reference: Data source reference

Optional:

• CDR3a, TRAV, TRAJ, TRAD, TRBD
• TRA_stitched, TRB_stitched
• Pseudo, MHC

Dependencies

• torch >= 2.0.0
• transformers >= 4.30.0
• numpy, pandas, tqdm
• python-Levenshtein
• nltk
• einops
• tidytcells >= 2.0.0
• mhcgnomes >= 1.8.0
• tcrpmhcdataset >= 0.2.0

License

Apache-2.0