zavavan/causal_reasoning_agent

1# File: app.py2 3import gradio as gr4from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig5from peft import PeftModel, PeftConfig6import torch7import regex as re8 9# Load PEFT adapter configuration10peft_config = PeftConfig.from_pretrained("unica/CLiMA")11 12# BitsAndBytes 4-bit config13bnb_config = BitsAndBytesConfig(14    load_in_4bit=True,15    bnb_4bit_quant_type="nf4",         # Most efficient for LLMs16    bnb_4bit_compute_dtype=torch.bfloat16,  # Use bfloat16 or float16 depending on your GPU17    bnb_4bit_use_double_quant=True18)19 20 21base_model = AutoModelForCausalLM.from_pretrained(22    peft_config.base_model_name_or_path,23    quantization_config=bnb_config,24    device_map="auto"25)26 27# Load adapter weights28model = PeftModel.from_pretrained(base_model, "unica/CLiMA")29 30# Load tokenizer31tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)32 33prompt_instruction_drug_reviews = f"""Given a drug review enclosed in triple quotes and a pair of entities E1 corresponding to the drug name and E2 corresponding to the treated condition, classify the relation holding between E1 and E2.34The relations are identified with 9 labels from 0 to 8. The meaning of the labels is the following:350 means that E1 causes E2361 means that E2 causes E1372 means that E1 enables E2383 means that E2 enables E1394 means that E1 prevents E2405 means that E2 prevents E1416 means that E1 hinders E2427 means that E2 hinders E1438 means that E1 and E2 are in a relation different than any of the previous ones.44Given X the label that you predicted, for the output use the format LABEL: X45"""46 47 48# Format prompt49def format_prompt(user_input, entity1, entity2):50#return f"Identify causal relations in the following clinical narrative:\n\n{user_input}\n\nEntity 1: {entity1}\nEntity 2: {entity2}\n\nCausal relations:"51  text = user_input52  prompt_text = f"Text:'''{text}'''"53  e1 = entity154  e2 = entity255  prompt_entities = f"\nEntities: E1: '''{e1}''', E2: '''{e2}'''"56  full_prompt = f"<USER> {prompt_instruction_drug_reviews} {prompt_text} {prompt_entities} <ASSISTANT>"57  return full_prompt58 59# Prediction function60def generate_relations(text, entity1, entity2):61    answer_label_regex_pattern =  re.compile(r'LABEL:?\s?(\d+)')62 63    64    prompt = format_prompt(text, entity1, entity2)65    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)66    outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)67    response = tokenizer.decode(outputs[0], skip_special_tokens=True)68    modelOut = response[len(prompt):].strip()  # remove prompt from output if echoed69    answer_match = answer_label_regex_pattern.search(modelOut)70    if answer_match:71        if answer_match.group(1)=='0':72            return f"""'{entity1}' causes  '{entity2}'"""73        elif answer_match.group(1)=='1':74            return f"""'{entity2}' causes  '{entity1}'"""75        elif answer_match.group(1)=='2':76            return f"""'{entity1}' enables  '{entity2}'"""77        elif answer_match.group(1)=='3':78            return f"""'{entity2}' enables  '{entity1}'"""79        elif answer_match.group(1)=='4':80            return f"""'{entity1}' prevents  '{entity2}'"""81        elif answer_match.group(1)=='5':82            return f"""'{entity2}' prevents  '{entity1}'"""83        elif answer_match.group(1)=='6':84            return f"""'{entity1}' hinders  '{entity2}'"""85        elif answer_match.group(1)=='7':86            return f"""'{entity2}' hinders  '{entity1}'"""87        elif answer_match.group(1)=='8':88            return f"""No causal relation between '{entity1}' and '{entity2}'"""89    else:90        return 'No causal relation could be extracted'91            92        93 94# Gradio UI95demo = gr.Interface(96    fn=generate_relations,97    inputs=[98        gr.Textbox(lines=10, label="Clinical Note or Drug Review Text"),99        gr.Textbox(label="Entity 1 (e.g., Drug)"),100        gr.Textbox(label="Entity 2 (e.g., Condition or Symptom)")101    ],102    outputs=gr.Textbox(label="Extracted Causal Relations"),103    title="Causal Relation Extractor with MedLlama",104    description="Paste your clinical note or drug review, and specify two target entities. This AI agent extracts drug-condition or symptom causal relations using a fine-tuned LLM adapter model.",105    examples=[106        ["Odynophagia: Was presumed due to mucositis from recent chemotherapy.", "chemotherapy", "mucositis"],107        ["patient's wife noticed erythema on patient's face. On [**3-27**]the visiting nurse [**First Name (Titles) 8706**][**Last Name (Titles)11282**]of a rash on his arms as well. The patient was noted to be febrile and was admitted to the [**Company 191**] Firm. In the EW, patient's Dilantin was discontinued and he was given Tegretol instead.", "Dilantin", "erythema on patient's face"],108        ["i had a urinary tract infection so bad that when i pee it smells but when i started taking ciprofloxacin it worked it’s a good medicine for a urinary tract infections.","ciprofloxacin","urinary tract infection"],109        ["when i first started using ziana, i only had acne in between my eyebrows, chin, and the nose area. my acne worsened while using it and then it got better. but after about 4 months of using it, it became ineffective. so i now have acne between my eyebrows, chin, cheeks, forehead, and the nose area. its great at first but after a while it made my face even worse than before i used the product.","ziana","acne"]110    ]111)112 113demo.launch()114