# 1. Import the required packages
import torch
import gradio as gr

from typing import Dict, Tuple
from transformers import pipeline

# 2. Define function to use our model on given text
def blood_request_classifier(text: str) -> Dict[str, float]:
    # Set up text classification pipeline
    classifier = pipeline(
        task="text-classification",
        model="AshenFdo/emergency_blood_request_classifier",
        device="cuda" if torch.cuda.is_available() else "cpu",
        top_k=None
    )
    outputs = classifier(text)[0]
    output_dict = {}
    for item in outputs:
        output_dict[item["label"]] = item["score"]

    return output_dict


# 3. Create a Gradio interface
description = """
A text classifier to determine whether a blood donation request is an **emergency** or **non-emergency**.

Fine-tuned from [DistilBERT](https://huggingface.co/distilbert/distilbert-base-uncased) on a 
[synthetic blood request urgency dataset](https://huggingface.co/datasets/AshenFdo/synthetic_blood_request_urgency_dataset).

See [source code on GitHub](https://github.com/AshenFdo/Blood-Request-Emergency-Classification-Model).
"""

demo = gr.Interface(fn=blood_request_classifier,
             inputs=gr.Textbox(
                 lines=4,
                 placeholder="Enter a blood request message here...",
                 label="Blood Request Text"
                 ),
             outputs=gr.Label(num_top_classes=2), # show top 2 classes (that's all we have)
             title="🩸 Emergency Blood Request Classifier",
             description=description,
             examples=[
                  ["Patient is in critical condition after surgery and urgently needs O- blood immediately or they may not survive."],
                  ["Hi, I am looking for a B+ blood donor for my father's scheduled knee replacement surgery next month."],
                  ["URGENT: Accident victim in ER needs AB+ blood NOW. Lives at stake, please respond immediately!"],
                  ["Our hospital is planning a blood donation camp next Saturday. All blood types welcome."],
                  ["A newborn baby in the ICU critically needs O+ blood within the next hour. Please help!"],
                ])

# 4. Launch the interface
if __name__ == "__main__":
    demo.launch()