html_class / app.py
SocialScrape's picture
Create app.py
506ca3a verified
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("Будь ласка, введіть текст для класифікації.")