Views
No views yet
1[Generate by superfamily] Superfamily=<xxx>
2or
3[Determine superfamily] Seq=<yyy>[Generate by superfamily] Superfamily=<Ankyrin repeat-containing domain superfamily>#You can also specify the first few amino acids of the protein sequence:
[Generate by superfamily] Superfamily=<Ankyrin repeat-containing domain superfamily> Seq=<MKRVL[Determine superfamily] Seq=<MAPGGMPREFPSFVRTLPEADLGYPALRGWVLQGERGCVLYWEAVTEVALPEHCHAECWGVVVDGRMELMVDGYTRVYTRGDLYVVPPQARHRARVFPGFRGVEHLSDPDLLPVRKR>1# you can replace the model_path with your local path
2CUDA_VISIBLE_DEVICES=0 python main.py --model "GreatCaptainNemo/ProLLaMA" --interactive
3# main.py is as follows 👇:1import argparse
2import json, os
3import torch
4from transformers import LlamaForCausalLM, LlamaTokenizer
5from transformers import GenerationConfig
6from tqdm import tqdm
7
8generation_config = GenerationConfig(
9 temperature=0.2,
10 top_k=40,
11 top_p=0.9,
12 do_sample=True,
13 num_beams=1,
14 repetition_penalty=1.2,
15 max_new_tokens=400
16)
17
18parser = argparse.ArgumentParser()
19parser.add_argument('--model', default=None, type=str,help="The local path of the model. If None, the model will be downloaded from HuggingFace")
20parser.add_argument('--interactive', action='store_true',help="If True, you can input instructions interactively. If False, the input instructions should be in the input_file.")
21parser.add_argument('--input_file', default=None, help="You can put all your input instructions in this file (one instruction per line).")
22parser.add_argument('--output_file', default=None, help="All the outputs will be saved in this file.")
23args = parser.parse_args()
24
25if __name__ == '__main__':
26 if args.interactive and args.input_file:
27 raise ValueError("interactive is True, but input_file is not None.")
28 if (not args.interactive) and (args.input_file is None):
29 raise ValueError("interactive is False, but input_file is None.")
30 if args.input_file and (args.output_file is None):
31 raise ValueError("input_file is not None, but output_file is None.")
32
33 load_type = torch.bfloat16
34 if torch.cuda.is_available():
35 device = torch.device(0)
36 else:
37 raise ValueError("No GPU available.")
38
39
40 model = LlamaForCausalLM.from_pretrained(
41 args.model,
42 torch_dtype=load_type,
43 low_cpu_mem_usage=True,
44 device_map='auto',
45 quantization_config=None
46 )
47 tokenizer = LlamaTokenizer.from_pretrained(args.model)
48
49 model.eval()
50 with torch.no_grad():
51 if args.interactive:
52 while True:
53 raw_input_text = input("Input:")
54 if len(raw_input_text.strip())==0:
55 break
56 input_text = raw_input_text
57 input_text = tokenizer(input_text,return_tensors="pt")
58
59 generation_output = model.generate(
60 input_ids = input_text["input_ids"].to(device),
61 attention_mask = input_text['attention_mask'].to(device),
62 eos_token_id=tokenizer.eos_token_id,
63 pad_token_id=tokenizer.pad_token_id,
64 generation_config = generation_config,
65 output_attentions=False
66 )
67 s = generation_output[0]
68 output = tokenizer.decode(s,skip_special_tokens=True)
69 print("Output:",output)
70 print("\n")
71 else:
72 outputs=[]
73 with open(args.input_file, 'r') as f:
74 examples =f.read().splitlines()
75 print("Start generating...")
76 for index, example in tqdm(enumerate(examples),total=len(examples)):
77 input_text = tokenizer(example,return_tensors="pt") #add_special_tokens=False ?
78
79 generation_output = model.generate(
80 input_ids = input_text["input_ids"].to(device),
81 attention_mask = input_text['attention_mask'].to(device),
82 eos_token_id=tokenizer.eos_token_id,
83 pad_token_id=tokenizer.pad_token_id,
84 generation_config = generation_config
85 )
86 s = generation_output[0]
87 output = tokenizer.decode(s,skip_special_tokens=True)
88 outputs.append(output)
89 with open(args.output_file,'w') as f:
90 f.write("\n".join(outputs))
91 print("All the outputs have been saved in",args.output_file)@article{lv2025prollama,
title={Prollama: A protein large language model for multi-task protein language processing},
author={Lv, Liuzhenghao and Lin, Zongying and Li, Hao and Liu, Yuyang and Cui, Jiaxi and Chen, Calvin Yu-Chian and Yuan, Li and Tian, Yonghong},
journal={IEEE Transactions on Artificial Intelligence},
year={2025},
publisher={IEEE}
}