Change to a faster setfit model

Change to an even faster, smaller model which uses embeddings instead. Less accurate but way faster for demo purposes.

app.py

@@ -1,13 +1,27 @@
 from fastapi import FastAPI
-from transformers import pipeline
-import os
+from transformers import AutoTokenizer, AutoModel
+import torch
+from sklearn.metrics.pairwise import cosine_similarity
 import logging
 
-# Set Hugging Face cache directory
-os.environ["HF_HOME"] = "/app/.cache"
-
+# Set up FastAPI app
 app = FastAPI()
-classifier = pipeline("zero-shot-classification", model="valhalla/distilbart-mnli-12-1")
+
+# Load tokenizer and model
+tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-small-en-v1")
+model = AutoModel.from_pretrained("BAAI/bge-small-en-v1")
+
+# Precompute embeddings for labels
+labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]
+label_embeddings = []
+
+for label in labels:
+    tokens = tokenizer(label, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        embedding = model(**tokens).last_hidden_state.mean(dim=1)
+    label_embeddings.append(embedding)
+
+label_embeddings = torch.vstack(label_embeddings)
 
 @app.get("/")
 async def root():
@@ -16,8 +30,17 @@ async def root():
 @app.post("/predict")
 async def predict(data: dict):
     logging.info(f"Received data: {data}")
-    labels = ["Mathematics", "Language Arts", "Social Studies", "Science"]
-    text = data["data"][0]  # Extract the text field properly
-    result = classifier(text, labels)
-    logging.info(f"Prediction result: {result}")
-    return {"label": result["labels"][0]}
+    text = data["data"][0]
+
+    # Compute embedding for input text
+    tokens = tokenizer(text, return_tensors="pt", padding=True, truncation=True)
+    with torch.no_grad():
+        text_embedding = model(**tokens).last_hidden_state.mean(dim=1)
+
+    # Compute cosine similarity
+    similarities = cosine_similarity(text_embedding, label_embeddings)[0]
+    best_label_idx = similarities.argmax()
+    best_label = labels[best_label_idx]
+
+    logging.info(f"Prediction result: {best_label}")
+    return {"label": best_label}