UCL-CSSB/PlasmidGPT

1import torch2from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig3 4device = 'cuda' if torch.cuda.is_available() else 'cpu'5print(f"Using device: {device}")6 7model = AutoModelForCausalLM.from_pretrained(8    "/Users/mcclainthiel/plasmidgpt-addgene-gpt2",9    trust_remote_code=True10).to(device)11model.eval()12 13tokenizer = AutoTokenizer.from_pretrained(14    "/Users/mcclainthiel/plasmidgpt-addgene-gpt2",15    trust_remote_code=True16)17 18start_sequence = 'ATGGCTAGCGAATTCGGCGCGCCT'19print(f"Start sequence: {start_sequence}\n")20 21input_ids = tokenizer.encode(start_sequence, return_tensors='pt').to(device)22 23outputs = model.generate(24    input_ids,25    max_length=300,26    num_return_sequences=1,27    temperature=1.0,28    do_sample=True,29    pad_token_id=tokenizer.pad_token_id,30    eos_token_id=tokenizer.eos_token_id31)32 33generated_sequence = tokenizer.decode(outputs[0], skip_special_tokens=True)34print(f"Generated sequence:\n{generated_sequence}\n")35print(f"Length: {len(generated_sequence)} bp")36