Update app.py

app.py

@@ -6,9 +6,22 @@ import threading
 import gradio as gr
 from sentence_transformers import SentenceTransformer, util
 
-EMBED_MODEL_ID = "sentence-transformers/all-MiniLM-L6-v2"
+# e5-small-v2 is trained specifically for asymmetric query->passage retrieval
+# (short query vs. longer descriptive text) using "query: "/"passage: "
+# prefixes. This handles short, title-style script lines ("Short Script on
+# AI.") far more reliably than a generic sentence-similarity model, which
+# has almost nothing to embed against in a 4-word input and falls back to
+# near-arbitrary rankings.
+EMBED_MODEL_ID = "intfloat/e5-small-v2"
 MAX_CONCURRENT_JOBS = 7
 
+# Below this cosine similarity, treat the pick as unreliable rather than
+# silently returning a low-confidence guess. e5 cosine scores for genuinely
+# relevant pairs typically sit higher than raw MiniLM scores did, but this
+# is a starting point — tune it against your own data if it flags too many
+# or too few scenes.
+LOW_CONFIDENCE_THRESHOLD = 0.78
+
 print("Loading embedding model...")
 embedder = SentenceTransformer(EMBED_MODEL_ID)
 print("✅ Ready")
@@ -73,15 +86,19 @@ def _clean_tags(tags: str) -> str:
     return ", ".join(kept) if kept else tags  # never return empty
 
 
-def _select_best_position(script: str, videos: list) -> int:
+def _select_best_position(script: str, videos: list):
     if not videos:
-        return None
+        return None, None
     tag_strings = [_clean_tags(v.get("tags", "")) for v in videos]
 
-    script_emb = embedder.encode(script, convert_to_tensor=True)
-    tag_embs = embedder.encode(tag_strings, convert_to_tensor=True)
+    # e5 convention: "query: " for the search text, "passage: " for the
+    # candidate documents being ranked against it.
+    query_emb = embedder.encode(f"query: {script}", convert_to_tensor=True)
+    passage_embs = embedder.encode(
+        [f"passage: {t}" for t in tag_strings], convert_to_tensor=True
+    )
 
-    sims = util.cos_sim(script_emb, tag_embs)[0]  # shape: (num_candidates,)
+    sims = util.cos_sim(query_emb, passage_embs)[0]  # shape: (num_candidates,)
     best_idx = int(sims.argmax())
     return videos[best_idx].get("position"), float(sims[best_idx])
 
@@ -116,6 +133,7 @@ def _select_best_videos(scenes_data: dict) -> list:
             "scene": scene_num,
             "script": script,
             "status": "ok",
+            "low_confidence": score < LOW_CONFIDENCE_THRESHOLD,
             "video": {
                 "position": chosen.get("position"),
                 "id": chosen.get("id"),
@@ -151,8 +169,9 @@ def gradio_select_best(json_text):
                 lines.append(f"| {r['scene']} | {r['script']} | — | — | *no candidates* | — |")
             else:
                 v = r["video"]
+                flag = " ⚠️" if r.get("low_confidence") else ""
                 lines.append(
-                    f"| {r['scene']} | {r['script']} | {v['id']} | {v['similarity']} | {v['tags']} | {v['pageURL']} |"
+                    f"| {r['scene']} | {r['script']} | {v['id']} | {v['similarity']}{flag} | {v['tags']} | {v['pageURL']} |"
                 )
         return "\n".join(lines)
     except Exception as e: