app.py

# File: app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer, BitsAndBytesConfig
from peft import PeftModel, PeftConfig
import torch

# Load PEFT adapter configuration
peft_config = PeftConfig.from_pretrained("unica/CLiMA")

# BitsAndBytes 4-bit config
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",         # Most efficient for LLMs
    bnb_4bit_compute_dtype=torch.bfloat16,  # Use bfloat16 or float16 depending on your GPU
    bnb_4bit_use_double_quant=True
)


base_model = AutoModelForCausalLM.from_pretrained(
    peft_config.base_model_name_or_path,
    quantization_config=bnb_config,
    device_map="auto"
)

# Load adapter weights
model = PeftModel.from_pretrained(base_model, "unica/CLiMA")

# Load tokenizer
tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)

# Format prompt
def format_prompt(user_input, entity1, entity2):
    return f"Identify causal relations in the following clinical narrative:\n\n{user_input}\n\nEntity 1: {entity1}\nEntity 2: {entity2}\n\nCausal relations:"

# Prediction function
def generate_relations(text, entity1, entity2):
    prompt = format_prompt(text, entity1, entity2)
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
    outputs = model.generate(**inputs, max_new_tokens=256, do_sample=False)
    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return response[len(prompt):].strip()  # remove prompt from output if echoed

# Gradio UI
demo = gr.Interface(
    fn=generate_relations,
    inputs=[
        gr.Textbox(lines=10, label="Clinical Note or Drug Review Text"),
        gr.Textbox(label="Entity 1 (e.g., Drug)"),
        gr.Textbox(label="Entity 2 (e.g., Condition or Symptom)")
    ],
    outputs=gr.Textbox(label="Extracted Causal Relations"),
    title="Causal Relation Extractor with MedLlama",
    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.",
    examples=[
        ["Odynophagia: Was presumed due to mucositis from recent chemotherapy.", "chemotherapy", "mucositis"],
        ["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"],
        ["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"],
        ["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"]
    ]
)

demo.launch()