Views
No views yet
pip install transformers1from transformers import AutoModelForCausalLM, AutoTokenizer
2import torch
3
4model_name = 'plant-dnamamba2-BPE'
5# load model and tokenizer
6model = AutoModelForCausalLM.from_pretrained(f'zhangtaolab/{model_name}', trust_remote_code=True)
7tokenizer = AutoTokenizer.from_pretrained(f'zhangtaolab/{model_name}', trust_remote_code=True)
8
9# example sequence and tokenization
10sequences = ['ATATACGGCCGNC','GGGTATCGCTTCCGAC']
11tokens = tokenizer(sequences,padding="longest")['input_ids']
12print(f"Tokenzied sequence: {tokenizer.batch_decode(tokens)}")
13
14# inference
15device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
16model.to(device)
17inputs = tokenizer(sequences, truncation=True, padding='max_length', max_length=512,
18 return_tensors="pt")
19inputs = {k: v.to(device) for k, v in inputs.items()}
20outs = model(
21 **inputs,
22 output_hidden_states=True
23)
24
25# get the final layer embeddings and prediction logits
26embeddings = outs['hidden_states'][-1].detach().numpy()
27logits = outs['logits'].detach().numpy()