Update app.py

app.py

@@ -3,7 +3,7 @@ import gradio as gr
 from datasets import load_dataset
 from sklearn.metrics.pairwise import cosine_similarity
 
-# -------- Config (safe defaults for CPU Spaces) --------
+# -------- Config --------
 SAMPLE_SIZE = int(os.getenv("SAMPLE_SIZE", "5000"))
 RANDOM_STATE = 42
 DEFAULT_INPUT = "I am so happy with this product"
@@ -18,7 +18,19 @@ def clean_text(text: str) -> str:
     text = re.sub(r"\s+", " ", text).strip()
     return text
 
+def _to_numpy(x):
+    try:
+        import torch
+        if hasattr(torch, "Tensor") and isinstance(x, torch.Tensor):
+            return x.detach().cpu().numpy()
+    except Exception:
+        pass
+    return np.asarray(x)
+
 def _l2norm(x: np.ndarray) -> np.ndarray:
+    x = x.astype(np.float32, copy=False)
+    if x.ndim == 1:
+        x = x.reshape(1, -1)
     return x / (np.linalg.norm(x, axis=1, keepdims=True) + 1e-12)
 
 # -------- Load sample data once --------
@@ -45,27 +57,31 @@ def load_generator():
     set_seed(RANDOM_STATE)
     return pipeline("text-generation", model="distilgpt2")
 
-# Map names → HF ids
+# HF model ids
 EMBEDDERS = {
     "MiniLM (fast)": "sentence-transformers/all-MiniLM-L6-v2",
     "MPNet (heavier)": "sentence-transformers/all-mpnet-base-v2",
     "DistilRoBERTa (paraphrase)": "sentence-transformers/paraphrase-distilroberta-base-v1",
 }
 
-# Cache for precomputed corpus embeddings per model
+# Cache for corpus embeddings per model
 _CORPUS_CACHE = {}
 
-def ensure_corpus_embeddings(model_name: str, texts: list[str]):
+def _encode_norm(model, texts):
+    """Encode with any sentence-transformers version and return (n,d) L2-normalized numpy."""
+    out = model.encode(texts, show_progress_bar=False)
+    out = _to_numpy(out)
+    return _l2norm(out)
+
+def ensure_corpus_embeddings(model_name: str, texts: list):
     if model_name in _CORPUS_CACHE:
         return _CORPUS_CACHE[model_name]
-    model_id = EMBEDDERS[model_name]
-    model = load_sentence_model(model_id)
-    emb = model.encode(texts, show_progress_bar=False, convert_to_numpy=True)
-    emb = _l2norm(emb)
+    model = load_sentence_model(EMBEDDERS[model_name])
+    emb = _encode_norm(model, texts)
     _CORPUS_CACHE[model_name] = emb
     return emb
 
-def top3_for_each_model(user_input: str, selected_models: list[str]):
+def top3_for_each_model(user_input: str, selected_models: list):
     df = load_sample_df()
     texts = df["clean_text"].tolist()
     rows = []
@@ -73,8 +89,7 @@ def top3_for_each_model(user_input: str, selected_models: list[str]):
         try:
             model = load_sentence_model(EMBEDDERS[name])
             corpus_emb = ensure_corpus_embeddings(name, texts)
-            q = model.encode([clean_text(user_input)], show_progress_bar=False, convert_to_numpy=True)
-            q = _l2norm(q)
+            q = _encode_norm(model, [clean_text(user_input)])
             sims = cosine_similarity(q, corpus_emb)[0]
             top_idx = sims.argsort()[-3:][::-1]
             for rank, i in enumerate(top_idx, start=1):
@@ -90,35 +105,40 @@ def top3_for_each_model(user_input: str, selected_models: list[str]):
                 "Model": name, "Rank": "-", "Similarity": "-",
                 "Tweet (clean)": f"[Error: {e}]", "Tweet (orig)": ""
             })
-    out = pd.DataFrame(rows, columns=["Model","Rank","Similarity","Tweet (clean)","Tweet (orig)"])
-    return out
+    return pd.DataFrame(rows, columns=["Model","Rank","Similarity","Tweet (clean)","Tweet (orig)"])
 
-def generate_and_pick_best(prompt: str, n_sequences: int, max_length: int, temperature: float, scorer_model_name: str):
+# -------- Generator + scoring (with progress) --------
+def generate_and_pick_best(prompt: str, n_sequences: int, max_length: int,
+                           temperature: float, scorer_model_name: str,
+                           progress=gr.Progress()):
+    progress(0.0, desc="Loading models…")
     gen = load_generator()
+    scorer = load_sentence_model(EMBEDDERS[scorer_model_name])
+
+    progress(0.3, desc="Generating candidates…")
     outputs = gen(
         prompt,
-        max_length=max_length,
-        num_return_sequences=n_sequences,
+        max_new_tokens=int(max_length),   # number of NEW tokens to generate
+        num_return_sequences=int(n_sequences),
         do_sample=True,
-        temperature=temperature,
-        pad_token_id=50256,  # silence warning
+        temperature=float(temperature),
+        pad_token_id=50256,
     )
     candidates = [o["generated_text"].strip() for o in outputs]
 
-    scorer_id = EMBEDDERS[scorer_model_name]
-    scorer = load_sentence_model(scorer_id)
-    q = scorer.encode([prompt], show_progress_bar=False, convert_to_numpy=True); q = _l2norm(q)
-    cand_vecs = scorer.encode(candidates, show_progress_bar=False, convert_to_numpy=True); cand_vecs = _l2norm(cand_vecs)
+    progress(0.7, desc="Scoring candidates…")
+    q = _encode_norm(scorer, [prompt])
+    cand_vecs = _encode_norm(scorer, candidates)
     sims = cosine_similarity(q, cand_vecs)[0]
     best_idx = int(sims.argmax())
+
     table = pd.DataFrame({
         "Rank": np.argsort(-sims) + 1,
         "Similarity": np.sort(sims)[::-1],
         "Generated Tweet": [c for _, c in sorted(zip(-sims, candidates))]
     })
-    best = candidates[best_idx]
-    best_score = float(sims[best_idx])
-    return best, best_score, table
+    progress(1.0)
+    return candidates[best_idx], float(sims[best_idx]), table
 
 # ---------------- UI ----------------
 with gr.Blocks(title="Sentiment140 Embeddings + Generation") as demo:
@@ -147,7 +167,7 @@ Type a tweet, get similar tweets from Sentiment140, and generate a new one.
 
     with gr.Row():
         n_seq = gr.Slider(3, 15, value=8, step=1, label="Number of candidates")
-        max_len = gr.Slider(30, 120, value=60, step=1, label="Max length")
+        max_len = gr.Slider(30, 120, value=60, step=1, label="Max length (new tokens)")
         temp = gr.Slider(0.5, 1.5, value=0.9, step=0.05, label="Temperature")
         scorer_model = gr.Dropdown(list(EMBEDDERS.keys()), value="MiniLM (fast)", label="Scorer embedding")