--- license: apache-2.0 base_model: - keiwoo/RoBERTpep - keiwoo/RoBERTtcr pipeline_tag: text-generation --- 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 [GRATCR dataset](https://doi.org/10.1109/JBHI.2024.3514089) 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-GRATCR' 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)
