Views
No views yet
facebook/esm2_t36_3B_UR50Dmeta-llama/Llama-3.1-8B-Instruct
1import torch
2from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer
3
4model = AutoModelForCausalLM.from_pretrained(
5 pretrained_model_name_or_path="xiao-fei/Prot2Text-V2-11B-Instruct-hf",
6 trust_remote_code=True,
7 torch_dtype=torch.bfloat16,
8 device_map="cuda"
9)
10
11esm_tokenizer = AutoTokenizer.from_pretrained("facebook/esm2_t36_3B_UR50D")
12llama_tokenizer = AutoTokenizer.from_pretrained(
13 pretrained_model_name_or_path="meta-llama/Llama-3.1-8B-Instruct",
14 pad_token='<|reserved_special_token_0|>'
15)
16
17example_sequence = (
18 "MCYSANGNTFLIVDNTQKRIPEEKKPDFVRENVGDLDGVIFVELVDGKYFMDYYNRDGSMAAFCGNGARAFSQ"
19 "YLIDRGWIKEKEFTFLSRAGEIKVIVDDSIWVRMPGVSEKKEMKVDGYEGYFVVVGVPHFVMEVKGIDELDVE"
20 "KLGRDLRYKTGANVDFYEVLPDRLKVRTYERGVERETKACGTGVTSVFVVYRDKTGAKEVKIQVPGGTLFLKE"
21 "ENGEIFLRGDVKRCSEE"
22)
23system_message = (
24 "You are a scientific assistant specialized in protein function "
25 "predictions. Given the sequence embeddings and other information "
26 "of a protein, describe its function clearly and concisely in "
27 "professional language. "
28)
29placeholder = '<|reserved_special_token_1|>'
30user_message = "Sequence embeddings: " + placeholder * (len(example_sequence)+2)
31tokenized_prompt = llama_tokenizer.apply_chat_template(
32 [
33 {"role": "system", "content": system_message},
34 {"role": "user", "content": user_message}
35 ],
36 add_generation_prompt=True,
37 tokenize=True,
38 return_tensors="pt",
39 return_dict=True
40)
41tokenized_sequence = esm_tokenizer(
42 example_sequence,
43 return_tensors="pt"
44)
45
46model.eval()
47generated = model.generate(
48 inputs=tokenized_prompt["input_ids"].to(model.device),
49 attention_mask=tokenized_prompt["attention_mask"].to(model.device),
50 protein_input_ids=tokenized_sequence["input_ids"].to(model.device),
51 protein_attention_mask=tokenized_sequence["attention_mask"].to(model.device),
52 max_new_tokens=1024,
53 eos_token_id=128009,
54 pad_token_id=128002,
55 return_dict_in_generate=False,
56 num_beams=4,
57 do_sample=False,
58)
59print(llama_tokenizer.decode(generated[0], skip_special_tokens=True))1@misc{prot2textv2,
2 title={Prot2Text-V2: Protein Function Prediction with Multimodal Contrastive Alignment},
3 author={Xiao Fei and Michail Chatzianastasis and Sarah Almeida Carneiro and Hadi Abdine and Lawrence P. Petalidis and Michalis Vazirgiannis},
4 year={2025},
5 eprint={2505.11194},
6 archivePrefix={arXiv},
7 primaryClass={cs.CE},
8 url={https://arxiv.org/abs/2505.11194},
9}