Upload 2 files

app.py

@@ -0,0 +1,49 @@
+import streamlit as st
+import torch
+from transformers import DistilBertForSequenceClassification, DistilBertTokenizerFast
+
+
+# Cache the model loading to speed up app restarts.
+@st.cache_resource
+def load_model_and_tokenizer():
+    model = DistilBertForSequenceClassification.from_pretrained(
+        "./results/checkpoint-1980"
+    )
+    tokenizer = DistilBertTokenizerFast.from_pretrained("distilbert-base-cased")
+    model.eval()
+    return model, tokenizer
+
+
+model, tokenizer = load_model_and_tokenizer()
+
+
+def classify_text(text: str) -> str:
+    """Tokenize the text and run inference."""
+    encoding = tokenizer(
+        text, return_tensors="pt", padding=True, truncation=True, max_length=128
+    )
+    with torch.no_grad():
+        outputs = model(**encoding)
+        logits = outputs.logits
+    predicted_class_id = torch.argmax(logits, dim=1).item()
+    id2label = model.config.id2label  # Assumes id2label was set during training.
+    predicted_label = (
+        id2label[predicted_class_id] if id2label else str(predicted_class_id)
+    )
+    return predicted_label
+
+
+# Build the Streamlit interface.
+st.title("Text Classification with DistilBERT")
+
+st.write("Enter text in the box below and click 'Classify' to see the predicted label.")
+
+# Text input area.
+user_text = st.text_area("Input Text", "")
+
+if st.button("Classify"):
+    if user_text.strip() == "":
+        st.error("Please enter some text to classify.")
+    else:
+        predicted_label = classify_text(user_text)
+        st.success(f"Predicted label: **{predicted_label}**")

requirements.txt

@@ -0,0 +1,3 @@
+streamlit
+transformers
+torch