Update app.py

app.py

@@ -2,49 +2,36 @@ import streamlit as st
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 import torch
 
-st.title("AI Model Text Classifier")
+st.title("Text Sentiment Classifier")
 
-# Use valid model names
+# Valid fine-tuned models
 model_list = [
-    ("roberta-base", "RoBERTa Base"),
-    ("bert-base-uncased", "BERT Base"),
-    ("distilbert-base-uncased", "DistilBERT Base")
+    ("distilbert-base-uncased-finetuned-sst-2-english", "DistilBERT (SST-2)"),
+    ("textattack/roberta-base-imdb", "RoBERTa (IMDB Sentiment)")
 ]
 
 @st.cache_resource
-def load_model(model_name):
-    return AutoModelForSequenceClassification.from_pretrained(model_name)
+def load_model_and_tokenizer(model_name):
+    model = AutoModelForSequenceClassification.from_pretrained(model_name)
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+    return model, tokenizer
 
-@st.cache_resource
-def load_tokenizer(model_name):
-    return AutoTokenizer.from_pretrained(model_name)
-
-# Load models + tokenizers
-models = [(load_model(name), load_tokenizer(name), label) for name, label in model_list]
+models = {label: load_model_and_tokenizer(name) for name, label in model_list}
 
-# User input
-text_input = st.text_area("Enter text to classify:")
+# UI
+text_input = st.text_area("Enter text:")
 
-# Choose model
-selected_model_label = st.selectbox("Select a model:", [label for _, _, label in models])
-
-# Find selected model
-for model, tokenizer, label in models:
-    if label == selected_model_label:
-        selected_model = model
-        selected_tokenizer = tokenizer
-        break
+model_choice = st.selectbox("Choose model:", list(models.keys()))
 
 if st.button("Classify"):
-    if text_input.strip() == "":
+    if not text_input.strip():
         st.warning("Please enter some text!")
     else:
-        inputs = selected_tokenizer(text_input, return_tensors="pt", truncation=True)
+        model, tokenizer = models[model_choice]
+        inputs = tokenizer(text_input, return_tensors="pt", truncation=True)
         with torch.no_grad():
-            outputs = selected_model(**inputs)
-        logits = outputs.logits
+            logits = model(**inputs).logits
         probs = torch.softmax(logits, dim=-1).squeeze().tolist()
-
-        st.write("### Classification probabilities:")
+        st.write("### Results:")
         for i, prob in enumerate(probs):
             st.write(f"Class {i}: {prob:.4f}")