import gradio as gr
from transformers import pipeline

# Load the token classification model
pipe = pipeline("token-classification", model="Clinical-AI-Apollo/Medical-NER", aggregation_strategy='simple')

def classify_text(text):
    # Get token classification results
    result = pipe(text)
    
    # Format the results to resemble the UI shown in the image
    formatted_output = ""
    for res in result:
        entity = res['entity_group']
        word = res['word']
        score = res['score']
        start = res['start']
        end = res['end']
        formatted_output += f"Entity: {entity}, Word: {word}, Score: {score:.4f}, Span: [{start}:{end}]\n"
    
    return formatted_output

# Gradio Interface
demo = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(lines=5, label="Enter Medical Text"),
    outputs=gr.Textbox(label="Entity Classification"),
    title="Medical Entity Classification",
    description="Enter medical-related text, and the model will classify medical entities.",
    examples=[
        ["45 year old woman diagnosed with CAD"],
        ["A 65-year-old male presents with acute chest pain and a history of hypertension."],
        ["The patient underwent a laparoscopic cholecystectomy."]
    ]
)

if __name__ == "__main__":
    demo.launch()