Update app.py

app.py

@@ -1,39 +1,54 @@
-import gradio as gr
-from transformers import pipeline
-
-from Gradio_UI import GradioUI
-
+# File: app.py
 
+import gradio as gr
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from peft import PeftModel, PeftConfig
+import torch
+
+# Load PEFT adapter configuration
+peft_config = PeftConfig.from_pretrained("unica/CLiMA")
+
+# Load base model
+base_model = AutoModelForCausalLM.from_pretrained(
+    peft_config.base_model_name_or_path,
+    torch_dtype=torch.float16,
+    device_map="auto"
+)
 
-# Load your fine-tuned model from Hugging Face Hub
-model = pipeline("text2text-generation", model='unica/CLiMA')  # Replace with your actual model repo name
-
+# Load adapter weights
+model = PeftModel.from_pretrained(base_model, "unica/CLiMA")
 
-# Define your prompt template (customize as needed)
-def format_prompt(user_input):
-    return f"Identify causal relations in the following clinical narrative:\n\n{user_input}\n\nCausal relations:"  # Modify if your model uses a different template
+# Load tokenizer
+tokenizer = AutoTokenizer.from_pretrained(peft_config.base_model_name_or_path)
 
-# Define prediction function
-def generate_relations(text):
-    prompt = format_prompt(text)
-    result = model(prompt, max_length=512, do_sample=False)
-    return result[0]['generated_text']
+# 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=512, do_sample=False)
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response[len(prompt):].strip()  # remove prompt from output if echoed
 
-# Gradio interface
+# Gradio UI
 demo = gr.Interface(
     fn=generate_relations,
-    inputs=gr.Textbox(lines=10, label="Clinical Note or Drug Review Text"),
+    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. This AI agent extracts drug-condition or symptom causal relations using a fine-tuned LLM.",
+    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=[
-        ["Patient reported severe headaches after starting amitriptyline."],
-        ["Lisinopril helped reduce the patient's blood pressure but caused persistent cough."],
-        ["After using Metformin, the patient experienced gastrointestinal discomfort."]
+        ["Patient reported severe headaches after starting amitriptyline.", "amitriptyline", "headaches"],
+        ["Lisinopril helped reduce the patient's blood pressure but caused persistent cough.", "Lisinopril", "cough"],
+        ["After using Metformin, the patient experienced gastrointestinal discomfort.", "Metformin", "gastrointestinal discomfort"]
     ]
 )
 
-# Launch the app
-demo.launch()
+demo.launch()