import os
import torch
import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification

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)
    probabilities = torch.softmax(outputs.logits, dim=-1)[0].detach().numpy()
    predicted_class = probabilities.argmax()

    return int(predicted_class), probabilities.tolist()

# Інтерфейс Streamlit
st.title("Longformer UK Text Classifier")
text_input = st.text_area("Введіть текст для класифікації")

if st.button("Класифікувати"):
    if text_input.strip() != "":
        predicted_class, probabilities = classify_text(text_input)
        st.markdown("### Результати класифікації:")
        st.write(f"**Клас:** {predicted_class}")
        st.write(f"**Ймовірності класів:** {probabilities}")
    else:
        st.warning("Будь ласка, введіть текст для класифікації.")