import gradio as gr
from transformers import pipeline

# Initialize the pipeline
pipe = pipeline(task="ner", 
                model='Mykes/rubert_ner_SDDCS', 
                tokenizer='Mykes/rubert_ner_SDDCS', 
                aggregation_strategy='first')

def process_text(text):
    # Convert input to lowercase as in your example
    results = pipe(text.lower())
    
    # Format the output
    output = []
    for entity in results:
        formatted_result = f"Type: {entity['entity_group']}\nWord: {entity['word']}\nScore: {entity['score']:.4f}\n"
        output.append(formatted_result)
    
    return "\n".join(output)

# Create Gradio interface
iface = gr.Interface(
    fn=process_text,
    inputs=gr.Textbox(lines=3, placeholder="Enter your text here..."),
    outputs=gr.Textbox(lines=10),
    title="Medical NER for Russian Text",
    description="This model identifies medical entities (diseases, symptoms, drugs, etc.) in Russian text.",
    examples=[
        ["У ребенка треога и норушения сна, потеря сознания, раньше ставили паническое расстройство. по назначению психиатра принимал атаракс без эффекта."],
        ["У женщины треога и норушения сна, потеря сознания, раньше ставили паническое расстройство. Подскажи хорошего психотервта в районе метро Октбрьской."],
        ["Моя дочка 14 лет и у нее начались месячные. Хотелось бы показать ее гиекологу женщине, живущей недалеко от мет шоссе энтузиастов?"],
        ["У меня болит живот, слабость, высокая температура. Что мне делать?"],
        ["У ребенка треога и норушения сна, потеря сознания, раньше ставили паническое расстройство. по назначению психиатра принимал атаракс без эффекта. Подскажи психиатра в районе метро Октбрьской."],
    ]
)

# Launch the interface
iface.launch()