README.md

---
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](https://proceedings.mlr.press/v311/zhang25a.html) 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](https://github.com/keiwoo/Pep2Tcr)