app.py

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

st.title("Text Sentiment Classifier")

# Valid fine-tuned models
model_list = [
    ("distilbert-base-uncased-finetuned-sst-2-english", "DistilBERT (SST-2)"),
    ("textattack/roberta-base-imdb", "RoBERTa (IMDB Sentiment)")
]

@st.cache_resource
def load_model_and_tokenizer(model_name):
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    return model, tokenizer

models = {label: load_model_and_tokenizer(name) for name, label in model_list}

# UI
text_input = st.text_area("Enter text:")

model_choice = st.selectbox("Choose model:", list(models.keys()))

if st.button("Classify"):
    if not text_input.strip():
        st.warning("Please enter some text!")
    else:
        model, tokenizer = models[model_choice]
        inputs = tokenizer(text_input, return_tensors="pt", truncation=True)
        with torch.no_grad():
            logits = model(**inputs).logits
        probs = torch.softmax(logits, dim=-1).squeeze().tolist()
        st.write("### Results:")
        for i, prob in enumerate(probs):
            st.write(f"Class {i}: {prob:.4f}")