import torch
import gradio as gr
from typing import Dict
from transformers import pipeline

def phishing_not_phishing_classifier(text:str) -> Dict[str, float]:
  phishing_not_phishing_classifier_pipeline = pipeline(task="text-classification",
                                                       model="twoweeksgetheart/learn_hf_phishing_not_phishing_text_classifier_rubert",
                                                       batch_size=32,
                                                       device="cuda" if torch.cuda.is_available() else "cpu",
                                                       top_k=None)
  outputs = phishing_not_phishing_classifier_pipeline(text)[0]

  outpid_dict = {}
  for item in outputs:
    outpid_dict[item["label"]] = item["score"]

  return outpid_dict

description = """
Является ли сообщение фишинговым?
"""

demo = gr.Interface(
    fn=phishing_not_phishing_classifier,
    inputs="text",
    outputs=gr.Label(num_top_classes=2),
    title="Классификация фишинговых сообщений",
    description=description,
)

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