Views
No views yet


1from models.progen.modeling_progen import ProGenForCausalLM
2import torch
3from tokenizers import Tokenizer
4import json
5
6# Define the model identifier from Hugging Face's model hub
7model_path = 'AntibodyGeneration/fine-tuned-progen2-small'
8
9# Load the model and tokenizer
10model = ProGenForCausalLM.from_pretrained(model_path)
11tokenizer = Tokenizer.from_file('tokenizer.json')
12
13# Define your sequence and other parameters
14target_sequence = 'MQIPQAPWPVVWAVLQLGWRPGWFLDSPDRPWNPPTFSPALLVVTEGDNATFTCSFSNTSESFVLNWYRMSPSNQTDKLAAFPEDRSQPGQDCRFRVTQLPNGRDFHMSVVRARRNDSGTYLCGAISLAPKAQIKESLRAELRVTERRAEVPTAHPSPSPRPAGQFQTLVVGVVGGLLGSLVLLVWVLAVICSRAARGTIGARRTGQPLKEDPSAVPVFSVDYGELDFQWREKTPEPPVPCVPEQTEYATIVFPSGMGTSSPARRGSADGPRSAQPLRPEDGHCSWPL'
15number_of_sequences = 2
16
17# Tokenize the sequence
18tokenized_sequence = tokenizer(target_sequence, return_tensors="pt")
19
20# Move model and tensors to CUDA if available
21device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
22model = model.to(device)
23tokenized_sequence = tokenized_sequence.to(device)
24
25# Generate sequences
26with torch.no_grad():
27 output = model.generate(**tokenized_sequence, max_length=1024, pad_token_id=tokenizer.pad_token_id, do_sample=True, top_p=0.9, temperature=0.8, num_return_sequences=number_of_sequences)
28
29# Decoding the output to get generated sequences
30generated_sequences = [tokenizer.decode(output_seq, skip_special_tokens=True) for output_seq in output]