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

# Замените 'username/имя-вашей-модели' на путь к вашей модели на Hugging Face
model_name = 'Yerzhxn/class_space'

# Загрузка токенизатора и модели
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Перемещение модели на устройство (если есть GPU)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# Интерфейс Streamlit
st.title("Тестирование классификации текста")
st.write("Введите текст, чтобы узнать предсказанный класс.")

# Поле ввода текста
input_text = st.text_area("Введите текст здесь", "")

if st.button("Предсказать"):
    if input_text:
        # Преобразование текста в формат, подходящий для модели
        inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
        inputs = {key: value.to(device) for key, value in inputs.items()}

        # Прогон текста через модель и получение предсказания
        with torch.no_grad():
            outputs = model(**inputs)

        # Преобразование выходных данных в вероятности и предсказание класса
        logits = outputs.logits
        predicted_class = torch.argmax(logits, dim=1).item()

        # Вывод результата
        st.write(f"Предсказанный класс: {predicted_class}")
    else:
        st.write("Пожалуйста, введите текст для классификации.")