app.py

# app.py
import os
import json
import time
import threading
import gradio as gr
from sentence_transformers import SentenceTransformer, util

# 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

# Gradio-level concurrency cap. This is kept well above MAX_CONCURRENT_JOBS
# so requests are handed to our function immediately instead of sitting in
# Gradio's own internal queue — capacity is enforced by our own counter
# (_active_jobs) below, which lets us return "server busy" the instant a
# request arrives rather than after it's been queued for a while.
GRADIO_CONCURRENCY_LIMIT = max(MAX_CONCURRENT_JOBS * 4, 20)

# 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")

# ══════════════════════════════════════════════════════════════════════════
# Concurrency / capacity tracking
# ══════════════════════════════════════════════════════════════════════════
_active_jobs_lock = threading.Lock()
_active_jobs = 0
_total_jobs_served = 0


def _job_started():
    global _active_jobs, _total_jobs_served
    with _active_jobs_lock:
        _active_jobs += 1
        _total_jobs_served += 1


def _job_finished():
    global _active_jobs
    with _active_jobs_lock:
        _active_jobs = max(0, _active_jobs - 1)


def _get_capacity_snapshot() -> dict:
    with _active_jobs_lock:
        active = _active_jobs
        total_served = _total_jobs_served
    return {
        "max_concurrent_jobs": MAX_CONCURRENT_JOBS,
        "active_jobs": active,
        "available_slots": max(0, MAX_CONCURRENT_JOBS - active),
        "total_jobs_served_since_startup": total_served,
    }


# job_id -> {"done": int, "total": int, "percent": float, "status": str}
# Lets a caller start a job with a job_id, then poll api_job_progress(job_id)
# from a second, concurrent request to get a live percentage while the
# first request is still running. Entries are pruned after being read once
# they're finished, so this stays small.
_job_progress_lock = threading.Lock()
_job_progress = {}


def _progress_set(job_id: str, done: int, total: int, status: str = "running"):
    if not job_id:
        return
    with _job_progress_lock:
        _job_progress[job_id] = {
            "done": done,
            "total": total,
            "percent": round((done / total * 100) if total else 100.0, 1),
            "status": status,
        }


def _progress_clear(job_id: str):
    if not job_id:
        return
    with _job_progress_lock:
        _job_progress.pop(job_id, None)


def _reserve_slot_or_busy():
    """
    Atomically checks capacity and reserves a slot in one step (so two
    requests can't both see a free slot and both get admitted). Returns
    None if a slot was reserved (caller MUST call _job_finished() when
    done, ideally in a finally block), or a capacity snapshot dict if the
    server is full and the caller should bail out immediately.
    """
    global _active_jobs, _total_jobs_served
    with _active_jobs_lock:
        if _active_jobs >= MAX_CONCURRENT_JOBS:
            return {
                "max_concurrent_jobs": MAX_CONCURRENT_JOBS,
                "active_jobs": _active_jobs,
                "available_slots": 0,
                "total_jobs_served_since_startup": _total_jobs_served,
            }
        _active_jobs += 1
        _total_jobs_served += 1
        return None


# ══════════════════════════════════════════════════════════════════════════
# Best-video selection via embedding similarity
#
# Instead of asking an LLM to pick a position number (which showed a strong
# bias toward always picking position 1, regardless of actual relevance),
# we embed the script line and every candidate's tag string, then pick the
# candidate with the highest cosine similarity. No generation involved, so
# there's no position bias and it's faster than even a 2-token LLM call —
# just one forward pass per script+candidate pair, batched.
# ══════════════════════════════════════════════════════════════════════════
# Pixabay boilerplate/metadata tags that describe the *asset* rather than
# its visual content, and can falsely spike similarity (e.g. "ai generated"
# just means the clip was made with generative AI tools — it says nothing
# about the subject matter — but it collides hard with any script mentioning
# "AI"). Stripped before embedding.
_NOISE_TAGS = {
    "ai generated", "live wallpaper", "background", "wallpaper",
    "beautiful wallpaper", "backdrop",
}


def _clean_tags(tags: str) -> str:
    parts = [t.strip() for t in tags.split(",")]
    kept = [t for t in parts if t.lower() not in _NOISE_TAGS]
    return ", ".join(kept) if kept else tags  # never return empty


def _select_best_position(script: str, videos: list):
    if not videos:
        return None, None
    tag_strings = [_clean_tags(v.get("tags", "")) for v in videos]

    # 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(query_emb, passage_embs)[0]  # shape: (num_candidates,)
    best_idx = int(sims.argmax())
    return videos[best_idx].get("position"), float(sims[best_idx])


def _select_best_videos(scenes_data: dict, progress_cb=None) -> list:
    """
    scenes_data: parsed JSON in the format:
      {"scenes": [{"scene": 1, "script": "...", "videos": [
          {"position": 1, "id": 847, "tags": "...", "pageURL": "...",
           "videoURL": "...", "thumbnailURL": "..."}, ...]}]}
    Returns a list of per-scene results with the chosen video's full info
    plus the similarity score that drove the choice.

    progress_cb, if given, is called after every scene as
    progress_cb(done, total) so callers (e.g. the Gradio UI) can render a
    real-time percentage-complete progress bar.
    """
    results = []
    scenes = scenes_data.get("scenes", [])
    total = len(scenes)

    for i, sc in enumerate(scenes):
        scene_num = sc.get("scene")
        script = sc.get("script", "")
        videos = sc.get("videos", [])

        if not videos:
            results.append({
                "scene": scene_num,
                "script": script,
                "status": "no_candidates",
                "video": None,
            })
        else:
            best_pos, score = _select_best_position(script, videos)
            chosen = next((v for v in videos if v.get("position") == best_pos), videos[0])
            results.append({
                "scene": scene_num,
                "script": script,
                "status": "ok",
                "low_confidence": score < LOW_CONFIDENCE_THRESHOLD,
                "video": {
                    "position": chosen.get("position"),
                    "id": chosen.get("id"),
                    "tags": chosen.get("tags"),
                    "pageURL": chosen.get("pageURL"),
                    "videoURL": chosen.get("videoURL"),
                    "thumbnailURL": chosen.get("thumbnailURL"),
                    "similarity": round(score, 4),
                },
            })

        if progress_cb is not None:
            progress_cb(i + 1, total)

    return results


# ══════════════════════════════════════════════════════════════════════════
# Gradio UI function
# ══════════════════════════════════════════════════════════════════════════
def gradio_select_best(json_text, progress=gr.Progress()):
    if not json_text or not json_text.strip():
        return "❌ Please paste the scenes JSON"

    # Reserve a slot up front. If the server is already at capacity this
    # returns immediately with a busy message — no queueing, no model call.
    busy_snapshot = _reserve_slot_or_busy()
    if busy_snapshot is not None:
        return (
            f"🚫 Server busy — all {busy_snapshot['max_concurrent_jobs']} "
            f"selection slots are in use. Please try again shortly."
        )

    try:
        try:
            data = json.loads(json_text)
        except json.JSONDecodeError as e:
            return f"❌ Invalid JSON: {e}"

        scenes_total = len(data.get("scenes", []))
        progress(0, desc=f"0/{scenes_total} scenes (0%)")

        def _on_progress(done, total):
            pct = int(done / total * 100) if total else 100
            progress(done / total if total else 1.0, desc=f"{done}/{total} scenes ({pct}%)")

        results = _select_best_videos(data, progress_cb=_on_progress)

        lines = ["| Scene | Script | Best Video ID | Similarity | Tags | Link |",
                 "|---|---|---|---|---|---|"]
        for r in results:
            if r["status"] == "no_candidates":
                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']}{flag} | {v['tags']} | {v['pageURL']} |"
                )
        return "\n".join(lines)
    except Exception as e:
        return f"❌ Error: {e}"
    finally:
        _job_finished()


# ══════════════════════════════════════════════════════════════════════════
# API functions — exposed via gr.api(), same reasoning as the whisper
# space: HF Spaces' proxy only recognizes Gradio's own route patterns, so
# custom FastAPI routes get silently swallowed. gr.api() registers through
# Gradio's own routing system instead.
# ══════════════════════════════════════════════════════════════════════════
def api_health() -> dict:
    """Health check. Returns service status."""
    return {"status": "ok", "model": EMBED_MODEL_ID, "device": "cpu"}


def api_status() -> dict:
    """
    Check server capacity — how many selection "channels"/slots are
    available right now.

    Returns:
      {
        "max_concurrent_jobs": 7,
        "active_jobs": 2,
        "available_slots": 5,
        "total_jobs_served_since_startup": 341
      }
    """
    return _get_capacity_snapshot()


def api_job_progress(job_id: str) -> dict:
    """
    Poll the live progress of a job started via api_select_best_videos with
    a matching job_id. Call this repeatedly from a second, concurrent
    request while the first request is still in flight.

    Returns:
      {"status": "running", "done": 3, "total": 10, "percent": 30.0}
      {"status": "not_found"}  — no job_id was ever started, or it's
                                  already been read once after finishing
                                  (entries are cleared after one read)
    """
    with _job_progress_lock:
        entry = _job_progress.pop(job_id, None) if job_id in _job_progress and _job_progress[job_id]["status"] == "done" else _job_progress.get(job_id)
    if entry is None:
        return {"status": "not_found"}
    return entry


def api_select_best_videos(scenes_json: str, job_id: str = "") -> dict:
    """
    Given a scenes JSON (script + candidate videos with tags per scene, e.g.
    your Pixabay search results), pick the single best-matching video per
    scene using embedding cosine similarity between the script line and each
    candidate's tags. No LLM generation involved, so there's no position
    bias and it runs fast even across many scenes/candidates.

    If all MAX_CONCURRENT_JOBS slots are currently in use, this returns
    {"status": "busy", ...} immediately rather than queueing the request.

    Real-time progress: pass a job_id (any unique string, e.g. a uuid), then
    poll /api_job_progress with that same job_id from a second concurrent
    request while this one is still running:

      import uuid, threading, time
      from gradio_client import Client
      client = Client("your-username/your-space")
      job_id = str(uuid.uuid4())

      def poll():
          while True:
              p = client.predict(job_id=job_id, api_name="/api_job_progress")
              print(p)  # {"status": "running", "done": 3, "total": 10, "percent": 30.0}
              if p.get("status") in ("done", "not_found"):
                  break
              time.sleep(0.5)

      threading.Thread(target=poll).start()
      result = client.predict(scenes_json=json.dumps(data), job_id=job_id,
                               api_name="/api_select_best_videos")

    Expected input shape:
      {"scenes": [{"scene": 1, "script": "...", "videos": [
          {"position": 1, "id": 847, "tags": "...", "pageURL": "...",
           "videoURL": "...", "thumbnailURL": "..."}, ...]}]}

    Returns:
      {"status": "ok", "results": [{"scene": 1, "script": "...", "status": "ok",
       "video": {"position":1,"id":847,"tags":"...","pageURL":"...",
       "videoURL":"...","thumbnailURL":"...","similarity":0.63}}], "took_seconds": 0.12}
      or {"status": "busy"/"error", "error": "..."}
    """
    if not scenes_json or not scenes_json.strip():
        return {"status": "error", "error": "Empty scenes JSON provided"}

    try:
        data = json.loads(scenes_json)
    except json.JSONDecodeError as e:
        return {"status": "error", "error": f"Invalid JSON: {e}"}

    # Reserve a slot up front so a burst of simultaneous requests can't all
    # slip past the capacity check before any of them increments it.
    busy_snapshot = _reserve_slot_or_busy()
    if busy_snapshot is not None:
        return {
            "status": "busy",
            "error": (
                f"Server busy: all {busy_snapshot['max_concurrent_jobs']} "
                f"selection slots are in use. Please try again shortly."
            ),
        }

    try:
        start_time = time.time()

        scenes_total = len(data.get("scenes", []))
        _progress_set(job_id, 0, scenes_total)

        def _on_progress(done, total):
            _progress_set(job_id, done, total)

        results = _select_best_videos(data, progress_cb=_on_progress)
        _progress_set(job_id, scenes_total, scenes_total, status="done")
        return {
            "status": "ok",
            "results": results,
            "took_seconds": round(time.time() - start_time, 3),
        }
    except Exception as e:
        _progress_clear(job_id)
        return {"status": "error", "error": f"Internal error: {e}"}
    finally:
        _job_finished()


with gr.Blocks(title="Best Video Selector — Embedding Similarity") as demo:
    gr.Markdown(f"""
    # 🎯 Best-Video Selector (Embedding Similarity)
    Paste your scenes JSON (script + Pixabay candidate videos with tags per
    scene). Each candidate is ranked by cosine similarity between the script
    line and its tags — no LLM generation, so there's no position bias and
    it's fast even across many scenes.
    Handles up to **{MAX_CONCURRENT_JOBS} simultaneous requests**; once all
    slots are busy, new requests get an instant "server busy" reply instead
    of waiting in a queue.

    **API users:** call via `gradio_client.Client("your-username/your-space")`,
    then `client.predict(scenes_json="...", api_name="/api_select_best_videos")`.
    Check `/api_status` first to see current server capacity.
    See the "View API" link in this page's footer for full request/response details.
    """)
    with gr.Row():
        with gr.Column():
            inp = gr.Textbox(
                label="Scenes JSON (script + candidate videos with tags)",
                placeholder='{"scenes": [{"scene": 1, "script": "...", "videos": [...]}]}',
                lines=20,
            )
            btn = gr.Button("🎯 Select Best Videos", variant="primary")
        with gr.Column():
            out = gr.Textbox(label="Result (markdown table)", lines=20)

    btn.click(
        fn=gradio_select_best,
        inputs=inp,
        outputs=out,
        concurrency_limit=GRADIO_CONCURRENCY_LIMIT,
    )

    gr.api(api_health, api_name="api_health")
    gr.api(api_status, api_name="api_status")
    gr.api(api_job_progress, api_name="api_job_progress")
    gr.api(
        api_select_best_videos,
        api_name="api_select_best_videos",
        concurrency_limit=GRADIO_CONCURRENCY_LIMIT,
    )


if __name__ == "__main__":
    demo.launch()