app.py

import gradio as gr
from transformers import pipeline


from Gradio_UI import GradioUI



final_answer = FinalAnswerTool()

# Load your fine-tuned model from Hugging Face Hub
model = pipeline("text2text-generation", model='unica/CLiMA')  # Replace with your actual model repo name


# 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

# 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']

    

# Gradio interface
demo = gr.Interface(
    fn=generate_relations,
    inputs=gr.Textbox(lines=10, label="Clinical Note or Drug Review Text"),
    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.",
    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."]
    ]
)

# Launch the app
demo.launch()