Initial probe deploy 2026-06-14 18:03:58

README.md

@@ -3,28 +3,32 @@ title: AnimoFlow Rewriter Probe
 emoji: 🔤
 colorFrom: blue
 colorTo: indigo
-sdk: docker
-app_port: 7860
+sdk: gradio
+sdk_version: 5.6.0
+app_file: app.py
 pinned: false
 short_description: Multilingual prompt rewriter probe (Qwen + RAFSL)
 hardware: zero-a10g
+suggested_hardware: zero-a10g
+preload_from_hub:
+  - "Qwen/Qwen2.5-1.5B-Instruct"
+  - "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2"
 ---
 
 # AnimoFlow Rewriter Probe
 
-Phase-0 verification of the multilingual prompt-rewriter front-end planned for AnimoFlow's text-to-motion pipeline. This Space loads Gemma-3-1B-it + a multilingual MiniLM retriever over the HumanML3D corpus, and rewrites arbitrary user input into HumanML3D-style English captions.
+Phase-0 verification of the multilingual prompt-rewriter front-end planned for AnimoFlow's text-to-motion pipeline. Loads a small instruction-tuned LLM + a multilingual MiniLM retriever over the HumanML3D corpus and rewrites arbitrary user input into HumanML3D-style English captions.
 
-- **Model:** [`google/gemma-3-1b-it`](https://huggingface.co/google/gemma-3-1b-it) (revision `dcc83ea841ab`)
-- **Retriever:** [`sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`](https://huggingface.co/sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2)
-- **Corpus:** HumanML3D captions (29K unique)
-- **Hardware:** CPU basic (free tier)
+- **Model:** `Qwen/Qwen2.5-1.5B-Instruct` by default (override via `MODEL_REPO` env / Space secret)
+- **Retriever:** `sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2`
+- **Corpus:** HumanML3D captions (52K unique)
+- **Hardware:** ZeroGPU (A10G burst, PRO-account free tier)
 
-## Endpoints
+## Usage
 
-- `POST /rewrite` — `{"prompt": "..."}` → `{"rewritten": "...", "latency_s": ..., ...}`
-- `GET /healthz` — readiness + cold-load metrics
-- `GET /ui` — Gradio sanity-check UI
+- **Gradio UI** — type a motion prompt in any language, press *Rewrite*.
+- **API** — `gradio_client.Client("AnimoFlow/rewriter-probe").predict(prompt, api_name="/rewrite")` returns `{rewritten, examples, latency_s, ...}`.
 
 ## Purpose
 
-This Space is a **measurement probe**, not a production service. It exists to verify the architecture decisions in the production rewriter plan before they land in `animoflow-web/nodes/prompt_rewriter/`.
+This Space is a **measurement probe** for the AnimoFlow rewriter plan, not a production service. Plan: `vault/wiki/concepts/Prompt rewriter front-end plan` in the AnimoFlow vault.

app.py

@@ -1,20 +1,26 @@
-"""HF Space probe: minimal /rewrite + healthcheck endpoints + Gradio UI."""
+"""HF Space probe: Gradio SDK + ZeroGPU.
+
+Exposes a /rewrite Gradio API endpoint that returns:
+  rewritten: str
+  examples: list[str]
+  latency_s: float
+  input_tokens: int
+  output_tokens: int
+  cold_load_s: float
+  uptime_s: float
+
+Plus a /healthz endpoint via gr.routes for the probe to poll.
+"""
 from __future__ import annotations
 
-import os
 import threading
 import time
-from contextlib import asynccontextmanager
 from pathlib import Path
 
 import gradio as gr
-import uvicorn
-from fastapi import FastAPI, HTTPException
-from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel, Field
 
 try:
-    import spaces  # HF ZeroGPU decorator
+    import spaces
     _HAS_SPACES = True
 except ImportError:
     _HAS_SPACES = False
@@ -22,13 +28,12 @@ except ImportError:
 from rewriter import Rewriter
 
 DATA_DIR = Path(__file__).parent / "data"
-PORT = int(os.environ.get("PORT", "7860"))
 
+_BOOT_AT = time.time()
 _REWRITER: Rewriter | None = None
 _LOAD_LOCK = threading.Lock()
-_LOAD_START_AT = time.time()
-_FIRST_INFERENCE_AT: float | None = None
 _LOAD_LATENCY: float | None = None
+_FIRST_INFERENCE_AT: float | None = None
 
 
 def get_rewriter() -> Rewriter:
@@ -42,70 +47,44 @@ def get_rewriter() -> Rewriter:
     return _REWRITER
 
 
-class RewriteRequest(BaseModel):
-    prompt: str = Field(..., min_length=1, max_length=2000)
-
-
-class RewriteResponse(BaseModel):
-    rewritten: str
-    examples: list[str]
-    latency_s: float
-    input_tokens: int
-    output_tokens: int
-
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    # Cold-load: instantiate the rewriter at boot so the first request doesn't pay the load cost.
-    print(f"[lifespan] starting cold-load at t+{time.time() - _LOAD_START_AT:.1f}s", flush=True)
-    get_rewriter()
-    print(f"[lifespan] cold-load done at t+{time.time() - _LOAD_START_AT:.1f}s", flush=True)
-    yield
-
-
-app = FastAPI(title="AnimoFlow Rewriter Probe", lifespan=lifespan)
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-
-@app.get("/healthz")
-def healthz():
-    ready = _REWRITER is not None
-    return {
-        "ready": ready,
-        "uptime_s": time.time() - _LOAD_START_AT,
-        "load_latency_s": _LOAD_LATENCY,
-        "first_inference_at_uptime_s": (_FIRST_INFERENCE_AT - _LOAD_START_AT) if _FIRST_INFERENCE_AT else None,
-        "device": _REWRITER.device if _REWRITER else None,
-    }
-
-
 def _do_rewrite(prompt: str) -> dict:
     return get_rewriter().rewrite(prompt)
 
 
-# On ZeroGPU, wrap the inference call so HF allocates a GPU burst for it.
+# On ZeroGPU, wrap inference so HF allocates a GPU burst.
 if _HAS_SPACES:
     _do_rewrite = spaces.GPU(duration=30)(_do_rewrite)
 
 
-@app.post("/rewrite", response_model=RewriteResponse)
-def rewrite(req: RewriteRequest):
+def rewrite_api(prompt: str) -> dict:
+    """Gradio API entry point exposed at /api/rewrite (and /rewrite for direct REST style)."""
     global _FIRST_INFERENCE_AT
-    try:
-        out = _do_rewrite(req.prompt)
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"rewriter failed: {type(e).__name__}: {e}")
+    if not prompt or not prompt.strip():
+        return {"rewritten": "", "examples": [], "latency_s": 0.0, "input_tokens": 0, "output_tokens": 0,
+                "cold_load_s": _LOAD_LATENCY or 0.0, "uptime_s": time.time() - _BOOT_AT, "error": "empty prompt"}
+    t0 = time.time()
+    out = _do_rewrite(prompt)
     if _FIRST_INFERENCE_AT is None:
         _FIRST_INFERENCE_AT = time.time()
-    return RewriteResponse(**out)
+    return {
+        **out,
+        "wall_latency_s": time.time() - t0,
+        "cold_load_s": _LOAD_LATENCY or 0.0,
+        "uptime_s": time.time() - _BOOT_AT,
+    }
+
+
+def healthz_api() -> dict:
+    return {
+        "ready": _REWRITER is not None,
+        "uptime_s": time.time() - _BOOT_AT,
+        "load_latency_s": _LOAD_LATENCY,
+        "first_inference_at_uptime_s": (_FIRST_INFERENCE_AT - _BOOT_AT) if _FIRST_INFERENCE_AT else None,
+        "device": _REWRITER.device if _REWRITER else None,
+        "model_id": _REWRITER.model_id if _REWRITER else None,
+    }
 
 
-# Gradio UI (mounted at /ui)
 def gradio_rewrite(prompt: str):
     if not prompt.strip():
         return "(empty)", "(empty)", 0.0
@@ -114,11 +93,17 @@ def gradio_rewrite(prompt: str):
     return out["rewritten"], examples_str, out["latency_s"]
 
 
+# Eager-load the rewriter at module import so the first Gradio call doesn't pay the load cost.
+# We do this AFTER the spaces decorator is wired up so ZeroGPU is initialised.
+get_rewriter()
+
+
 with gr.Blocks(title="AnimoFlow Rewriter Probe") as demo:
-    gr.Markdown("# AnimoFlow Rewriter Probe — Gemma-3-1B-it + RAFSL")
+    gr.Markdown("# AnimoFlow Rewriter Probe — Qwen2.5-1.5B-Instruct + RAFSL on ZeroGPU")
     gr.Markdown(
         "Type a motion prompt in any language. The rewriter normalises it to a "
-        "HumanML3D-style English caption."
+        "HumanML3D-style English caption. Powered by Qwen2.5-1.5B-Instruct + a multilingual "
+        "MiniLM retriever over the 52K HumanML3D caption corpus."
     )
     with gr.Row():
         with gr.Column():
@@ -128,10 +113,20 @@ with gr.Blocks(title="AnimoFlow Rewriter Probe") as demo:
             out_text = gr.Textbox(label="Rewritten (HumanML3D-style English)", lines=2)
             out_examples = gr.Textbox(label="Retrieved exemplars", lines=4)
             out_latency = gr.Number(label="Latency (s)", precision=3)
-    btn.click(gradio_rewrite, inputs=[inp], outputs=[out_text, out_examples, out_latency])
+    btn.click(gradio_rewrite, inputs=[inp], outputs=[out_text, out_examples, out_latency], api_name="rewrite_ui")
+
+    # Pure API endpoint with JSON dict response — what the probe + clients should hit.
+    with gr.Row(visible=False):
+        api_in = gr.Textbox()
+        api_out = gr.JSON()
+        api_btn = gr.Button()
+        api_btn.click(rewrite_api, inputs=[api_in], outputs=[api_out], api_name="rewrite")
 
-app = gr.mount_gradio_app(app, demo, path="/ui")
+    with gr.Row(visible=False):
+        hz_btn = gr.Button()
+        hz_out = gr.JSON()
+        hz_btn.click(healthz_api, inputs=[], outputs=[hz_out], api_name="healthz")
 
 
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=PORT)
+    demo.queue().launch(server_name="0.0.0.0", server_port=7860)