BERTic / app.py
ngamulin's picture
Create app.py
1f131a2 verified
Raw
History Blame Contribute Delete
1.5 kB
import torch
import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from huggingface_hub import hf_hub_download
REPO_ID = "SentilyticsOPJ/bertic_model"
FILENAME = "bertic_model.pt"
path = hf_hub_download(repo_id=REPO_ID, filename=FILENAME)
obj = torch.load(path, map_location="cpu", weights_only=False)
hf_model_name = obj["hf_model_name"]
labels = obj["label_classes"]
max_len = obj["max_len"]
tokenizer = AutoTokenizer.from_pretrained(hf_model_name)
model = AutoModelForSequenceClassification.from_pretrained(
hf_model_name,
num_labels=len(labels),
)
model.load_state_dict(obj["state_dict"])
model.eval()
def predict(text):
if not text.strip():
return ""
enc = tokenizer(
text,
max_length=max_len,
padding="max_length",
truncation=True,
return_tensors="pt",
)
with torch.no_grad():
logits = model(
input_ids=enc["input_ids"],
attention_mask=enc["attention_mask"],
).logits
return labels[int(logits.argmax(1))]
demo = gr.Interface(
fn=predict,
inputs=gr.Textbox(lines=3, placeholder="Unesi tekst...", label="Text"),
outputs=gr.Textbox(label="Sentiment"),
title="Sentiment Analysis (BERTić, fine-tuned)",
description="Five-class sentiment: mixed, negative, neutral, positive, sarcastic.",
examples=["Volim kavu", "Ovaj doktor je loš", "Dan je bio ok"],
)
demo.launch()