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

MODEL_NAME = "SocialScrape/longformer-my-classifier"

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)
    probs = torch.softmax(outputs.logits, dim=-1)[0].tolist()
    pred_class = int(torch.argmax(outputs.logits))
    return {
        "Predicted Class": pred_class,
        "Class Probabilities": probs
    }

# Інтерфейс лише для API (без UI)
iface = gr.Interface(
    fn=classify_text,
    inputs=gr.Textbox(),
    outputs="json",
    live=False
)

iface.launch(show_api=True, share=False, inbrowser=False)