import gradio as gr
from transformers import pipeline

classifier = pipeline("text-classification", model="havocy28/VetBERTDx")

def predict(text):
    if not text or not text.strip():
        return "Wpisz objawy zwierzaka"
    
    result = classifier(text.strip())[0]
    raw_label = result["label"].upper()
    score = result["score"]
    
    # Mapujemy prawdziwe etykiety Twojego modelu
    label_map = {
        "UNKNOWN": "HIGH",   # <-- to jest Twój przypadek
        "LOW": "LOW",
        "MEDIUM": "MEDIUM",
        "HIGH": "HIGH"
    }
    
    final_label = label_map.get(raw_label, raw_label)
    
    return f"{final_label} – pewność: {score:.0%}"

gr.Interface(
    fn=predict,
    inputs=gr.Textbox(lines=6, label="Opis objawów", placeholder="Pies ma czerwone dziąsła, śmierdzi z pyska..."),
    outputs=gr.Textbox(label="Ocena ryzyka VetBERTDx"),
    title="VetBERTDx – triage weterynaryjny",
    description="Działa! Model zwrócił UNKNOWN → mapujemy na HIGH"
).launch()