import torch
from transformers import AutoTokenizer, AutoProcessor, AutoModelForCausalLM
from PIL import Image
REPO = "dnagpt/OmniGene-4-MM-merged"
tok = AutoTokenizer.from_pretrained(REPO)
proc = AutoProcessor.from_pretrained(REPO)
model = AutoModelForCausalLM.from_pretrained(
REPO, torch_dtype=torch.bfloat16, device_map="auto",
)
model.eval()
img = Image.open("molecule.png").convert("RGB")
msgs = [{"role": "user", "content": [
{"type": "image"},
{"type": "text", "text": "Please list the functional groups of the molecule."},
]}]
text = proc.apply_chat_template(msgs, add_generation_prompt=True, tokenize=False)
inp = proc(text=text, images=[img], return_tensors="pt").to(model.device)
out = model.generate(**inp, max_new_tokens=160, do_sample=False)
print(tok.decode(out[0][inp.input_ids.shape[1]:], skip_special_tokens=True))
prompt = '''### Instruction:
Determine if the two sequences below are structurally related (like paraphrases).
### Sequence 1:
MSRIGNKVIVLPAGVELANNDNVVTVKGPKGELTREFSKDIEIRVEGTEVTLHRPNDSKEMKTIHGTTRALL
### Sequence 2:
MSRIGNKVIVLPAGVELANNDNVVTVKGPKGELTREFSKDIEIRVEGTEVTLHRPNDSKEMKTIHGTTRALL
### Answer:
'''
ids = tok(prompt, return_tensors="pt").input_ids.to(model.device)
out = model.generate(ids, max_new_tokens=8, do_sample=False)
print(tok.decode(out[0][ids.shape[1]:], skip_special_tokens=True))