Views
No views yet
1from transformers import AutoTokenizer, AutoConfig, AutoModelForSequenceClassification
2
3
4# configuration dict for different TRAFICA versions
5config_dict = {
6 'TRAFICA (BPE-1)': {'model_path':'Allanxu/TRAFICA-BPE1',
7 'tokenizer_path':'Allanxu/TRAFICA-BPE1',
8 'tokenization':'BPE'
9 },
10 'TRAFICA (BPE-2)': {'model_path':'Allanxu/TRAFICA-BPE2',
11 'tokenizer_path':'zhihan1996/DNABERT-2-117M',
12 'tokenization':'BPE'
13 },
14 'TRAFICA (4-mer)': {'model_path':'Allanxu/TRAFICA-4_mer',
15 'tokenizer_path':'Allanxu/TRAFICA-4_mer',
16 'tokenization':'4_mer'
17 },
18 'TRAFICA (5-mer)': {'model_path':'Allanxu/TRAFICA-5_mer',
19 'tokenizer_path':'Allanxu/TRAFICA-5_mer',
20 'tokenization':'5_mer'
21 },
22 'TRAFICA (6-mer)': {'model_path':'Allanxu/TRAFICA-6_mer',
23 'tokenizer_path':'Allanxu/TRAFICA-6_mer',
24 'tokenization':'6_mer'
25 },
26 'TRAFICA (base-level)': {'model_path':'Allanxu/TRAFICA-Base_level',
27 'tokenizer_path':'Allanxu/TRAFICA-Base_level',
28 'tokenization':'Base-level'
29 }
30}
31
32# tokenizer
33Tokenizer = AutoTokenizer.from_pretrained(config_dict['TRAFICA (base-level)']['tokenizer_path'], trust_remote_code=True)
34
35# model
36config = AutoConfig.from_pretrained(config_dict['TRAFICA (base-level)']['model_path'])
37config.num_labels = 1
38
39
40model = AutoModelForSequenceClassification.from_pretrained(config_dict['TRAFICA (base-level)']['model_path'], config=config, trust_remote_code=True)1from peft import PeftModel
2import torch
3
4lora_path = '/<Path of fine-tuned LoRA>/Base-level/PRJEB3289/10000/ATF7_TGGGCG30NCGT' # example for TF ATF7
5
6# LoRA and Affinity predictor
7state_dict = torch.load(os.path.join(lora_path,"predict_head_weights.pth"), weights_only=True)
8model.classifier.load_state_dict( state_dict['PREDICT_HEAD'] )
9model = PeftModel.from_pretrained(model, os.path.join(lora_path,"lora_adapter"))lora_path by your local path1from util.py import piece_sequences # Src/util.py
2
3# Input construction
4sequences = ['CCAGAAGACAACTTGTAGAAATAAGCAAAA', 'ATTGCGCCCCAGCCCCACACCCACACGCAT']
5tokens_batch = piece_sequences(sequences, config_dict['TRAFICA (base-level)']['tokenization'])
6# tokens_batch = ['C C A G A A G A C A A C T T G T A G A A A T A A G C A A A A', 'A T T G C G C C C C A G C C C C A C A C C C A C A C G C A T']
7inputs = Tokenizer(tokens_batch, return_tensors="pt", padding=True)
8
9# Prediction
10with torch.no_grad():
11 outputs = model(**inputs)
12
13logit = outputs.logits
14print(f"Predicted relative affinities: {logit.flatten()}")