Views
No views yet
1import torch
2from transformers import AutoTokenizer, AutoModelForCausalLM
3from peft import PeftModel
4
5# 1. Load model
6tokenizer = AutoTokenizer.from_pretrained("DUTIR-BioNLP/RexDrug-base", trust_remote_code=True)
7model = AutoModelForCausalLM.from_pretrained(
8 "dlutIR/RexDrug-base",
9 torch_dtype=torch.bfloat16,
10 device_map="auto",
11 trust_remote_code=True,
12)
13model = PeftModel.from_pretrained(model, "DUTIR-BioNLP/RexDrug-adapter")
14model.eval()
15
16# 2. Prepare input
17messages = [
18 {"role": "system", "content": "You are an expert in biomedical drug-drug relation extraction. ..."},
19 {"role": "user", "content": "Target sentence: ... \nContext paragraph: ..."},
20]
21input_text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
22inputs = tokenizer(input_text, return_tensors="pt").to(model.device)
23
24# 3. Generate
25with torch.no_grad():
26 outputs = model.generate(**inputs, max_new_tokens=1024, do_sample=False)
27response = tokenizer.decode(outputs[0][inputs["input_ids"].shape[-1]:], skip_special_tokens=True)
28print(response)