File size: 1,166 Bytes
03e7ee5 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
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",
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()
|