Views
No views yet
1from transformers import AutoTokenizer, AutoModelForCausalLM
2tokenizer = AutoTokenizer.from_pretrained('quidangz/LLamaEE-8B-Instruct-ZeroShot')
3model = AutoModelForCausalLM.from_pretrained(
4 'quidangz/LLamaEE-8B-Instruct-ZeroShot',
5 torch_dtype="auto",
6 device_map="cuda",
7)
8
9if tokenizer.pad_token is None:
10 tokenizer.pad_token = tokenizer.eos_token
11 model.config.pad_token_id = model.config.eos_token_id
12
13user_prompt = """
14 Extract events and their components from text **strictly using ONLY the provided Event List** below and **MUST** strictly adhere to the output format.
15 Format output as '<event_type>: <trigger_word> | <role1>: <argument1> | <role2>: <argument2>' and separate multiple events with '|'. Return 'None' if no events are identified.
16 Event List: {ee_labels}
17 Text: {text}
18"""
19
20query = 'Drug-induced fever due to diltiazem'
21ee_labels = ["adverse event", "potential therapeutic event", "Treatment.Drug", "Treatment.Route,Subject.Disorder", "Treatment.Freq", "Treatment", "Effect"]
22
23user_prompt = user_prompt.format(ee_labels=ee_labels, text=query)
24
25messages = [
26 {
27 "role": "system",
28 "content": "You are an expert in Event Extraction (EE) task."
29 },
30 {
31 "role": "user",
32 "content": user_prompt
33 }
34]
35
36text = tokenizer.apply_chat_template(
37 messages,
38 tokenize=False,
39 add_generation_prompt=True
40 )
41
42model_inputs = tokenizer(text, return_tensors="pt").to(model.device)
43
44generated_ids = model.generate(
45 **model_inputs,
46 max_new_tokens=512,
47)
48
49generated_ids = [
50 output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)
51]
52
53response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]
54
55print(response) # adverse event: due to | Effect: fever | Treatment: diltiazem | Treatment.Drug: diltiazem1@misc{LlamaEE-8B-Instruct-ZeroShot,
2 title={LlamaEE: An Large Language Model for Relation Extraction},
3 author={Qui Dang Ba},
4 year={2025},
5 publisher={Huggingface},
6}