Update app.py

app.py

@@ -1,21 +1,36 @@
 import gradio as gr
-from transformers import pipeline
+from setfit import SetFitModel
+import numpy as np
 
+# 👇 use your exact model repo
 MODEL_ID = "DelaliScratchwerk/text-period-setfit"
-pipe = pipeline("text-classification", model=MODEL_ID, return_all_scores=True)
 
-def predict(txt):
-    scores = pipe(txt)[0]
-    scores = sorted(scores, key=lambda d: d["score"], reverse=True)
-    top = scores[0]["label"]
-    table = {s["label"]: round(float(s["score"]), 4) for s in scores}
-    return top, table
+# The label order must match the order you used during training
+LABELS = ["pre-1900","1900–1945","1946–1990","1991–2008","2009–2015","2016–2018","2019–2022","2023–present"]
+
+model = SetFitModel.from_pretrained(MODEL_ID)
+
+def predict(txt: str):
+    # SetFit expects a list; returns probabilities per label
+    probs = model.predict_proba([txt])[0]              # shape: (num_labels,)
+    order = np.argsort(probs)[::-1]                    # descending
+    top_label = LABELS[order[0]]
+    table = {LABELS[i]: float(probs[i]) for i in order}
+    return top_label, table
+
+examples = [
+    "Schools went remote during the pandemic; everyone wore N95s and used Zoom.",
+    "Sputnik launched and kicked off the space race.",
+    "MySpace was the most popular social network for a while.",
+    "TikTok creators exploded in popularity.",
+]
 
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Textbox(lines=8, label="Paste text"),
     outputs=[gr.Label(label="Predicted Period"), gr.JSON(label="Scores")],
-    title="Text → Time Period"
+    title="Text → Time Period (SetFit)",
+    examples=examples,
 )
 
 if __name__ == "__main__":