metadata
license: apache-2.0
pipeline_tag: text-generation
base_model:
- keiwoo/RoBERTpep
- keiwoo/RoBERTtcr
This repository contains relevant information of our research article 'Leveraging contrastive learning for conditional and robust TCR sequence generation'.
Transformers
You can use Pep2Tcr trained on TCRGen dataset with Transformers. To get started, install the necessary dependencies to setup your environment:
pip install -U transformers torch
2. Run the generation
from transformers import EncoderDecoderModel, AutoTokenizer
import torch
torch.cuda.manual_seed_all(42)
device = torch.device('cuda:1')
model_name = 'keiwoo/Pep2Tcr-TCRGen'
model = EncoderDecoderModel.from_pretrained(
model_name,
dtype=torch.float16, # If your device does not support `BF16` precision mode
# (GPU based on NVIDIA's Ampere and its subsequent architectures), you should comment this line.
).to(device)
model.eval()
tokenizer = AutoTokenizer.from_pretrained(model_name, do_lower_case=False)
sample_kwargs = {
'do_sample': True,
'top_k': 9,
'top_p': 0.93,
'temperature': 0.9,
'num_return_sequences': 16,
'repetition_penalty': 1.1,
'no_repeat_ngram_size': 3,
'max_new_tokens': 24,
}
beam_kwargs = {
'num_beams': 16,
'num_return_sequences': 16,
'repetition_penalty': 1.1,
'no_repeat_ngram_size': 3,
'max_new_tokens': 24,
}
inputs = "AAGIGILTV"
tokenized = tokenizer(" ".join(inputs), return_tensors="pt", ).to(device)
outputs = model.generate(**tokenized, **sample_kwargs) # If you want different sequences every generation, uncomment this line.
# outputs = model.generate(**tokenized, **beam_kwargs) # If you want same sequences every generation (most possible sequence the model think), uncomment this line.
tcr = tokenizer.batch_decode(outputs, skip_special_tokens=True)
tcr = [i.replace(' ', '') for i in tcr]
print(tcr)
# ['CASSQEGLAGGTDTQYF', 'CASSQEGLAGGYNEQFF', 'CASSQDPGLAGGYNEQFF', ...]
3. Reproduce the model
Please refer to https://github.com/keiwoo/Pep2Tcr