Upload app.py with huggingface_hub

app.py

@@ -1,18 +1,18 @@
-"""Vaaani Flagship — CPU inference API for the free Hugging Face Space.
-
-Loads the Qwen2.5-3B base GGUF + the curriculum LoRA (downloaded from an HF model
-repo at first request) and serves three shapes:
-    GET  /                      health (returns immediately so the Space boots fast)
-    POST /chat                  {messages:[...]} -> {reply}
-    POST /chat/stream           Server-Sent Events, token-by-token (perceived speed)
-    POST /v1/chat/completions   OpenAI-compatible drop-in for the existing gateway
-
-Design notes for the free tier (2 vCPU / 16 GB):
-  * Model loads LAZILY on the first chat call, NOT at boot — so the health check
-    passes instantly and the Space goes "Running" without a startup timeout.
-  * n_threads = 2 (both free vCPUs); n_ctx kept small (curriculum prompts are short).
-  * Default GGUF is Q4_K_M (~2 GB) — faster on CPU than Q8 at ~no quality loss for
-    this rigidly-scripted task. Override with env vars to serve Q8 on a paid Space.
+"""Vaaani engine — CPU, OpenAI-compatible inference for the Vaaani RAG backend.
+
+This is a HEADLESS model engine (no standalone UI). It is the drop-in for the RAG
+product's VAAANI_LLM_BASE_URL, and it routes by the request's `model` field:
+
+    vaaani-base      -> base Qwen2.5-3B GGUF, NO adapter   (general RAG / chat / ingest)
+    vaaani-flagship  -> base + curriculum LoRA              (the Root-Bridge tutor)
+
+The "no symbol before Grade 5" firewall is applied to sub-G5 tutor output only
+(detected from "Grade N" in the system prompt); general RAG and G5 are untouched.
+
+Endpoints:
+    GET  /                     health JSON
+    POST /v1/chat/completions  OpenAI-compatible, model-routed  (the RAG calls this)
+    POST /chat, /chat/stream   simple test helpers
 """
 import os
 import re
@@ -20,36 +20,40 @@ import json
 from typing import List, Optional
 
 from fastapi import FastAPI
-from fastapi.responses import HTMLResponse, JSONResponse, StreamingResponse
+from fastapi.responses import JSONResponse, StreamingResponse
 from pydantic import BaseModel
 from huggingface_hub import hf_hub_download
 
-HERE = os.path.dirname(os.path.abspath(__file__))
-
-# ── config (all overridable via Space "Variables" — no secrets in code) ──────
-REPO      = os.environ.get("VAAANI_MODEL_REPO", "Shaankar39/vaaani-flagship-gguf")
-BASE_FILE = os.environ.get("VAAANI_BASE_FILE", "vaaani-base-q4_k_m.gguf")
-LORA_FILE = os.environ.get("VAAANI_LORA_FILE", "vaaani-flagship-lora-f16.gguf")
+# ── config (override via Space Variables; secrets like HF_TOKEN stay secrets) ──
+REPO          = os.environ.get("VAAANI_MODEL_REPO", "Shaankar39/vaaani-flagship-gguf")
+BASE_FILE     = os.environ.get("VAAANI_BASE_FILE", "vaaani-base-q4_k_m.gguf")
+LORA_FILE     = os.environ.get("VAAANI_LORA_FILE", "vaaani-flagship-lora-f16.gguf")
+BASE_NAME     = os.environ.get("VAAANI_LLM_MODEL", "vaaani-base")
+FLAGSHIP_NAME = os.environ.get("VAAANI_FLAGSHIP_MODEL", "vaaani-flagship")
 N_THREADS = int(os.environ.get("N_THREADS", "2"))
 N_CTX     = int(os.environ.get("N_CTX", "2048"))
 N_BATCH   = int(os.environ.get("N_BATCH", "256"))
 MAX_TOK   = int(os.environ.get("MAX_TOKENS", "512"))
 
-app = FastAPI(title="Vaaani Flagship (CPU)")
-_llm = None  # lazy singleton
+app = FastAPI(title="Vaaani Engine (CPU)")
+_llms = {}  # "base" | "flagship" -> Llama (lazy)
 
 
-def get_llm():
-    global _llm
-    if _llm is None:
+def get_llm(use_lora: bool):
+    key = "flagship" if use_lora else "base"
+    if key not in _llms:
         from llama_cpp import Llama
-        base = hf_hub_download(REPO, BASE_FILE)
-        lora = hf_hub_download(REPO, LORA_FILE)
-        _llm = Llama(
-            model_path=base, lora_path=lora,
-            n_ctx=N_CTX, n_threads=N_THREADS, n_batch=N_BATCH, verbose=False,
-        )
-    return _llm
+        kw = dict(model_path=hf_hub_download(REPO, BASE_FILE),
+                  n_ctx=N_CTX, n_threads=N_THREADS, n_batch=N_BATCH, verbose=False)
+        if use_lora:
+            kw["lora_path"] = hf_hub_download(REPO, LORA_FILE)
+        _llms[key] = Llama(**kw)
+    return _llms[key]
+
+
+def _use_lora(model_name: Optional[str]) -> bool:
+    """Apply the curriculum adapter only when the flagship model is requested."""
+    return (model_name or "").strip() == FLAGSHIP_NAME
 
 
 class Msg(BaseModel):
@@ -59,6 +63,7 @@ class Msg(BaseModel):
 
 class ChatReq(BaseModel):
     messages: List[Msg]
+    model: Optional[str] = None
     temperature: float = 0.2
     max_tokens: Optional[int] = None
     stream: bool = False
@@ -69,10 +74,6 @@ def _msgs(req: ChatReq):
 
 
 # ── "no symbol before Grade 5" firewall (deterministic, serving-layer) ───────
-# The model gets this right ~6/7 of the time, but a hard rule needs a hard guarantee.
-# For any sub-G5 lesson we strip slash-phoneme notation (/b/, /th/) and IPA chars from
-# the OUTPUT — the same scrub sound_lessons.strip_symbols() applies to the training data.
-# G5 is allowed to reveal symbols, so the firewall never touches it.
 _SLASH_PHONEME = re.compile(r"/[A-Za-zθðŋʃʒʧʤ]+/")
 _IPA_CHARS = re.compile(r"[θðŋʃʒʧʤæɪʊəɔɑːʰˈˌ]")
 _GRADE_RE = re.compile(r"Grade\s+(\d+)")
@@ -99,29 +100,27 @@ def _firewall(text: str) -> str:
 
 
 class StreamScrubber:
-    """Streaming-safe version of the firewall for sub-G5. Lets G1–G4 stream
-    token-by-token (so it feels alive at ~2.5 tok/s) while still removing
-    /phoneme/ slash-notation and IPA chars across token boundaries. When a '/'
-    arrives we hold text until we can tell a phoneme (/b/) from a real slash
-    (on/off); a run longer than a phoneme is flushed as-is."""
+    """Streaming-safe firewall for sub-G5: scrubs /phoneme/ and IPA across token
+    boundaries while still streaming token-by-token. Holds text after a '/' until
+    it can tell a phoneme (/b/) from a real slash (on/off)."""
 
     def __init__(self):
-        self.hold = ""  # buffered text from an unresolved '/'
+        self.hold = ""
 
     def feed(self, s: str) -> str:
         out = []
         for ch in s:
             if self.hold:
-                if ch == "/":                      # closing slash -> it was a phoneme
+                if ch == "/":
                     out.append("that sound")
                     self.hold = ""
                 elif (ch.isalpha() or ch in "θðŋʃʒʧʤ") and len(self.hold) <= 6:
-                    self.hold += ch                # still possibly a phoneme
-                else:                              # not a phoneme -> flush held + this char
+                    self.hold += ch
+                else:
                     out.append(self.hold + ch)
                     self.hold = ""
             elif ch == "/":
-                self.hold = "/"                    # start holding
+                self.hold = "/"
             else:
                 out.append(ch)
         return _IPA_CHARS.sub("", "".join(out))
@@ -131,123 +130,16 @@ class StreamScrubber:
         return _IPA_CHARS.sub("", rest)
 
 
-@app.get("/health")
+@app.get("/")
 def health():
-    return {"status": "ok", "model": BASE_FILE, "lora": LORA_FILE,
-            "loaded": _llm is not None}
-
-
-@app.get("/lessons")
-def lessons():
-    with open(os.path.join(HERE, "lessons.json"), encoding="utf-8") as f:
-        return JSONResponse(json.load(f))
-
-
-@app.get("/", response_class=HTMLResponse)
-def home():
-    return UI_HTML
-
-
-UI_HTML = """<!doctype html>
-<html lang="en"><head>
-<meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1">
-<title>Vaaani — learn English by discovery</title>
-<style>
-  :root{--ink:#1a2238;--bg:#f6f7fb;--card:#fff;--accent:#4f46e5;--accent2:#0e9f6e;--muted:#6b7280}
-  *{box-sizing:border-box}
-  body{margin:0;font-family:system-ui,-apple-system,Segoe UI,Roboto,sans-serif;background:var(--bg);color:var(--ink)}
-  header{padding:20px 24px;background:linear-gradient(120deg,#4f46e5,#0e9f6e);color:#fff}
-  header h1{margin:0;font-size:22px;letter-spacing:.3px}
-  header p{margin:4px 0 0;opacity:.9;font-size:13px}
-  main{max-width:820px;margin:0 auto;padding:18px}
-  .lessons{display:flex;flex-wrap:wrap;gap:8px;margin-bottom:14px}
-  .lessons button{border:1px solid #d7dae5;background:var(--card);color:var(--ink);
-    padding:8px 12px;border-radius:999px;cursor:pointer;font-size:13px}
-  .lessons button:hover{border-color:var(--accent);color:var(--accent)}
-  #chat{background:var(--card);border:1px solid #e6e8f0;border-radius:14px;min-height:320px;
-    padding:14px;overflow-y:auto;max-height:60vh}
-  .bubble{padding:10px 13px;border-radius:12px;margin:8px 0;max-width:88%;white-space:pre-wrap;line-height:1.5}
-  .user{background:#eef0ff;margin-left:auto;color:#312e81}
-  .bot{background:#f0faf5;border:1px solid #d6f0e3}
-  .meta{font-size:11px;color:var(--muted);margin:2px 4px}
-  form{display:flex;gap:8px;margin-top:12px}
-  input{flex:1;padding:11px 13px;border:1px solid #d7dae5;border-radius:10px;font-size:15px}
-  .send{background:var(--accent);color:#fff;border:none;padding:0 18px;border-radius:10px;cursor:pointer;font-size:15px}
-  .send:disabled{opacity:.5;cursor:default}
-  .hint{font-size:12px;color:var(--muted);margin-top:8px}
-</style></head>
-<body>
-<header>
-  <h1>🌉 Vaaani</h1>
-  <p>English sounds &amp; word-roots, taught by discovery — runs fully on CPU.</p>
-</header>
-<main>
-  <div class="lessons" id="lessons"></div>
-  <div id="chat"></div>
-  <form id="f"><input id="q" placeholder="Pick a lesson above, or type to the child's tutor…" autocomplete="off"><button class="send" id="send">Send</button></form>
-  <div class="hint">Free CPU tier: the first reply is slow (model loads once); after that ~a minute per lesson, streaming as it thinks.</div>
-</main>
-<script>
-let messages = [];           // full chat incl. system
-const chat = document.getElementById('chat');
-const q = document.getElementById('q');
-const sendBtn = document.getElementById('send');
-
-function add(role, text){
-  const b = document.createElement('div');
-  b.className = 'bubble ' + (role === 'user' ? 'user' : 'bot');
-  b.textContent = text;
-  chat.appendChild(b); chat.scrollTop = chat.scrollHeight; return b;
-}
-
-async function stream(){
-  sendBtn.disabled = true;
-  const bot = add('bot', '…');
-  let acc = '';
-  try{
-    const res = await fetch('/chat/stream', {method:'POST',headers:{'Content-Type':'application/json'},
-      body: JSON.stringify({messages, max_tokens: 220})});
-    const reader = res.body.getReader(); const dec = new TextDecoder(); let buf='';
-    while(true){
-      const {value, done} = await reader.read(); if(done) break;
-      buf += dec.decode(value, {stream:true});
-      let i;
-      while((i = buf.indexOf('\\n\\n')) >= 0){
-        const line = buf.slice(0, i).trim(); buf = buf.slice(i+2);
-        if(!line.startsWith('data:')) continue;
-        const payload = line.slice(5).trim();
-        if(payload === '[DONE]') continue;
-        try{ const d = JSON.parse(payload).delta; if(d){ acc += d; bot.textContent = acc; chat.scrollTop = chat.scrollHeight; } }catch(e){}
-      }
-    }
-  }catch(e){ bot.textContent = 'Error: ' + e; }
-  if(acc) messages.push({role:'assistant', content: acc});
-  sendBtn.disabled = false; q.focus();
-}
-
-document.getElementById('f').addEventListener('submit', e=>{
-  e.preventDefault(); const t = q.value.trim(); if(!t) return;
-  if(messages.length === 0){ messages.push({role:'system', content:'You are Vaaani, a warm, encouraging tutor for a child learning English. Teach by asking and playing. Tiny warm sentences, emoji.'}); }
-  messages.push({role:'user', content:t}); add('user', t); q.value=''; stream();
-});
-
-fetch('/lessons').then(r=>r.json()).then(ls=>{
-  const box = document.getElementById('lessons');
-  ls.forEach(l=>{
-    const btn = document.createElement('button'); btn.textContent = l.label;
-    btn.onclick = ()=>{ chat.innerHTML=''; messages=[{role:'system',content:l.system},{role:'user',content:l.user}];
-      add('user', l.user); stream(); };
-    box.appendChild(btn);
-  });
-});
-</script>
-</body></html>"""
+    return {"status": "ok", "engine": "vaaani", "base": BASE_FILE, "lora": LORA_FILE,
+            "models": [BASE_NAME, FLAGSHIP_NAME], "loaded": list(_llms.keys())}
 
 
 @app.post("/chat")
 def chat(req: ChatReq):
     sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
-    out = get_llm().create_chat_completion(
+    out = get_llm(_use_lora(req.model)).create_chat_completion(
         messages=_msgs(req), temperature=req.temperature,
         max_tokens=req.max_tokens or MAX_TOK, stream=False)
     reply = out["choices"][0]["message"]["content"]
@@ -258,11 +150,8 @@ def chat(req: ChatReq):
 
 @app.post("/chat/stream")
 def chat_stream(req: ChatReq):
-    """SSE: emit {'delta': '<token>'} chunks, then [DONE]. The child sees words appear
-    as they generate. For sub-G5 the firewall must see the full text, so we generate
-    fully then emit the scrubbed result (turns are short — still fast); G5 streams live."""
     sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
-    llm = get_llm()
+    llm = get_llm(_use_lora(req.model))
 
     def gen():
         scrub = StreamScrubber() if sub_g5 else None
@@ -285,23 +174,25 @@ def chat_stream(req: ChatReq):
 
 @app.post("/v1/chat/completions")
 def openai_compat(req: ChatReq):
-    """OpenAI-shaped, so the existing gateway can point here unchanged."""
+    """OpenAI-compatible — the Vaaani RAG backend points VAAANI_LLM_BASE_URL here and
+    calls /v1/chat/completions with model=vaaani-base or vaaani-flagship."""
     sub_g5 = _is_sub_g5(_system_text(_msgs(req)))
-    llm = get_llm()
+    llm = get_llm(_use_lora(req.model))
     if req.stream:
         def gen():
-            if sub_g5:
-                out = llm.create_chat_completion(
+            scrub = StreamScrubber() if sub_g5 else None
+            for chunk in llm.create_chat_completion(
                     messages=_msgs(req), temperature=req.temperature,
-                    max_tokens=req.max_tokens or MAX_TOK, stream=False)
-                out["choices"][0]["message"]["content"] = _firewall(
-                    out["choices"][0]["message"]["content"])
-                yield f"data: {json.dumps(out)}\n\n"
-            else:
-                for chunk in llm.create_chat_completion(
-                        messages=_msgs(req), temperature=req.temperature,
-                        max_tokens=req.max_tokens or MAX_TOK, stream=True):
-                    yield f"data: {json.dumps(chunk)}\n\n"
+                    max_tokens=req.max_tokens or MAX_TOK, stream=True):
+                if sub_g5:
+                    delta = chunk["choices"][0]["delta"].get("content")
+                    if delta:
+                        chunk["choices"][0]["delta"]["content"] = scrub.feed(delta)
+                yield f"data: {json.dumps(chunk)}\n\n"
+            if scrub:
+                tail = scrub.flush()
+                if tail:
+                    yield f"data: {json.dumps({'choices':[{'index':0,'delta':{'content':tail},'finish_reason':None}]})}\n\n"
             yield "data: [DONE]\n\n"
         return StreamingResponse(gen(), media_type="text/event-stream")
     out = llm.create_chat_completion(