import gradio as gr
import numpy as np
from optimum.onnxruntime import ORTModelForSequenceClassification
from transformers import AutoTokenizer

MODEL_REPO = "Ashybalka/xlm-roberta-taxonomy-healthcare-de"
TOP_K = 5

print("Loading model...")
model = ORTModelForSequenceClassification.from_pretrained(
    MODEL_REPO,
    subfolder="onnx-int8",
    file_name="model_quantized.onnx",
)
tokenizer = AutoTokenizer.from_pretrained(MODEL_REPO)
id2label = model.config.id2label
print(f"Loaded. {len(id2label)} classes.")


def classify(title: str, description: str):
    title = (title or "").strip()
    description = (description or "").strip()

    if not title and not description:
        return {}, "Enter the job title and/or vacancy description."

    if title and description:
        text = f"{title} [SEP] {description}"
    else:
        text = title or description

    inputs = tokenizer(
        text,
        return_tensors="pt",
        truncation=True,
        max_length=512,
        padding=True,
    )

    outputs = model(**inputs)
    logits = outputs.logits[0]

    logits_np = logits.numpy() if hasattr(logits, "numpy") else np.asarray(logits)
    exp = np.exp(logits_np - logits_np.max())
    probs = exp / exp.sum()

    scores = {id2label[i]: float(probs[i]) for i in range(len(probs))}
    top_label = max(scores, key=scores.get)
    confidence = scores[top_label]

    summary = f"**{top_label}** ({confidence:.1%})"
    return scores, summary


with gr.Blocks(title="Healthcare Classifier (DE)") as demo:
    gr.Markdown(
"""
# Healthcare-Klassifikator für deutsche Stellenanzeigen

Fine-tuned **XLM-RoBERTa** (INT8 quantized, ONNX) — classifies German-language healthcare and medical job postings into **11 healthcare subcategories**: Medical Doctors, Nursing & Care, Therapy & Rehabilitation, Psychology & Psychotherapy, Pharmacy & Pharma Industry, Medical Technology & Devices, Laboratory & Diagnostics, Dentistry, Emergency & Paramedic Services, Medical Assistance, Healthcare Management & Administration.

**Test metrics:** Accuracy 94.5% · F1 macro 87.8% · F1 weighted 94.5%

Designed to be used downstream of the [main industry classifier](https://huggingface.co/Ashybalka/xlm-roberta-taxonomy-main-de) — route only listings tagged as `Healthcare & Medical` into this model for fine-grained role classification. Labeled via 3-model LLM consensus with soft-label training.

> **Note:** Strong performance on the top 7 categories (F1 ≥ 0.90). Smaller classes like `Psychology & Psychotherapy` (F1 0.65) and `Healthcare Management & Administration` (F1 0.75) have ambiguous boundaries in training data — use with care.

[Model card →](https://huggingface.co/Ashybalka/xlm-roberta-taxonomy-healthcare-de)
"""
    )

    with gr.Row():
        with gr.Column(scale=1):
            title_in = gr.Textbox(
                label="Job Title",
                placeholder="z.B. Gesundheits- und Krankenpfleger",
                lines=1,
            )
            desc_in = gr.Textbox(
                label="Job Description",
                placeholder="z.B. Examinierte Pflegefachkraft für die Betreuung von Bewohnern in unserer Senioreneinrichtung...",
                lines=10,
            )
            btn = gr.Button("Klassifizieren", variant="primary")

        with gr.Column(scale=1):
            top_out = gr.Markdown(label="Top prediction")
            scores_out = gr.Label(
                num_top_classes=TOP_K,
                label="Healthcare subcategories",
            )

    gr.Examples(
        examples=[
            [
                "Facharzt für Innere Medizin",
                "Für unsere internistische Abteilung suchen wir einen Facharzt (m/w/d) für Innere Medizin. Abgeschlossene Facharztausbildung, Approbation und Berufserfahrung in der stationären Versorgung erforderlich. Mitwirkung am Bereitschaftsdienst.",
            ],
            [
                "Gesundheits- und Krankenpfleger (m/w/d) Intensivstation",
                "Examinierte Pflegefachkraft für unsere Intensivstation. Durchführung der Grund- und Behandlungspflege, Überwachung intensivpflichtiger Patienten, Dokumentation. Abgeschlossene Ausbildung erforderlich, Fachweiterbildung Intensivpflege wünschenswert. Schichtdienst.",
            ],
            [
                "Physiotherapeut",
                "Physiotherapeut (m/w/d) für unsere ambulante Praxis. Behandlung von Patienten im Bereich Orthopädie, Neurologie und Sportphysiotherapie. Abgeschlossene Ausbildung als Physiotherapeut erforderlich, Fortbildungen in manueller Therapie oder Lymphdrainage von Vorteil.",
            ],
            [
                "Notfallsanitäter",
                "Notfallsanitäter (m/w/d) für unseren Rettungsdienst. Einsatz im qualifizierten Krankentransport und in der Notfallrettung. Abgeschlossene Ausbildung zum Notfallsanitäter, Führerschein Klasse C1 erforderlich. Bereitschaft zum Schichtdienst.",
            ],
            [
                "Zahnmedizinische Fachangestellte (ZFA)",
                "ZFA für unsere moderne Zahnarztpraxis. Stuhlassistenz, Patientenbetreuung, Terminmanagement, Abrechnung. Abgeschlossene Ausbildung als Zahnmedizinische Fachangestellte erforderlich. Erfahrung mit DAMPSOFT oder Charly von Vorteil.",
            ],
            [
                "Apotheker (m/w/d)",
                "Approbierter Apotheker für unsere Filialapotheke in Berlin. Beratung von Kunden, Rezepturherstellung, Sicherstellung der Arzneimittelversorgung. Abgeschlossenes Pharmaziestudium und Approbation erforderlich. Eigenständige Arbeitsweise und Kundenorientierung.",
            ],
            [
                "MTLA Zentrallabor",
                "Medizinisch-technische Laboratoriumsassistenz für unser Zentrallabor. Durchführung von hämatologischen, klinisch-chemischen und mikrobiologischen Analysen. Abgeschlossene Ausbildung als MTLA erforderlich. Bereitschaft zum Schichtdienst und Wochenenddienst.",
            ],
            [
                "Medizintechniker Krankenhaus",
                "Medizintechniker (m/w/d) für unser Klinikum. Wartung, Reparatur und Prüfung medizintechnischer Geräte (MTK/STK), Beratung der Anwender, Dokumentation nach MPBetreibV. Abgeschlossene Ausbildung im Bereich Medizintechnik oder Elektrotechnik erforderlich.",
            ],
        ],
        inputs=[title_in, desc_in],
        label="Beispiele",
    )

    btn.click(
        fn=classify,
        inputs=[title_in, desc_in],
        outputs=[scores_out, top_out],
    )

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