import gradio as gr
from transformers import BertTokenizer, BertForSequenceClassification
import torch

model = BertForSequenceClassification.from_pretrained("DmitriySv/ticket_classifier")
tokenizer = BertTokenizer.from_pretrained("DmitriySv/ticket_classifier")

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = model.to(device)
model.eval()

def classify(text):
    inputs = tokenizer(text, padding=True, truncation=True, max_length=512, return_tensors="pt").to(device)
    
    with torch.no_grad():
        logits_task1, logits_task2 = model(**inputs)

    pred_task1 = torch.argmax(logits_task1, dim=1).item()
    pred_task2 = torch.argmax(logits_task2, dim=1).item()

    return {"Тип": pred_task1, "Приоритет": pred_task2}

interface = gr.Interface(
    fn=classify,  
    inputs=gr.Textbox(label="Введите запрос для классификации"),
    outputs=[gr.Label(label="Тип"), gr.Label(label="Приоритет")],
    title="Классификация запроса по типу и приоритету",
    description="Классификация запроса по типу и приоритету."
)

interface.launch()