import os
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr

MODEL_NAME = "SocialScrape/longformer-my-classifier"

# hf_token = os.getenv("HF_TOKEN")
# print("HF_TOKEN:", hf_token)

tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)

# Основна функція класифікації
def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True)
    outputs = model(**inputs)
    probabilities = torch.softmax(outputs.logits, dim=-1)[0].detach().numpy()
    predicted_class = probabilities.argmax()

    return {
        "Клас": int(predicted_class),
        "Ймовірності класів": probabilities.tolist()
    }

# Градіо інтерфейс
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(label="Введіть текст для класифікації"),
    outputs=gr.JSON(label="Результати класифікації"),
    title="Longformer UK Text Classifier"
)

# Обов'язково запускати
iface.launch()