html_classifier / app.py
SocialScrape's picture
Update app.py
987bdff verified
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()