v3: swap to Llama 4 Scout via Together-hosted endpoint (82% task, 72% Governance)

app.py

@@ -1,9 +1,15 @@
-"""HF Space demo for construction-code-cite (v1.1 LoRA).
-
-Uses transformers + peft (not mlx-lm) so it runs on Space Linux CPU.
-Fetches the OSHA 1926 corpus at cold boot from the public HF dataset,
-loads the LoRA adapter on top of Qwen 2.5 1.5B-Instruct, and serves
-strict-JSON hazard + citation predictions through Gradio.
+"""HF Space demo for construction-code-cite (v3 · Llama 4 Scout 17B-16E).
+
+The v3 model is 109B total params (17B active MoE) and does not fit in-Space.
+Inference goes to Together AI's hosted endpoint; the Space runs the RAG
+pipeline + verifier + Gradio UI only. If TOGETHER_API_KEY is missing or the
+endpoint returns an error, we fall back to the v2 (Llama 3.2 3B) in-Space
+adapter so the demo never goes dark.
+
+Set the following secrets in the Space:
+- TOGETHER_API_KEY   (required for v3)
+- V3_MODEL_ID        (default: rigidhat/llama-4-scout-17b-construction-codecite-v3)
+- FALLBACK_ADAPTER   (default: rigidhat/llama-3.2-3b-construction-codecite-v2)
 """
 from __future__ import annotations
 
@@ -15,8 +21,14 @@ from pathlib import Path
 
 import gradio as gr
 
-BASE_MODEL = os.environ.get("BASE_MODEL", "meta-llama/Llama-3.2-3B-Instruct")
-ADAPTER_REPO = os.environ.get("ADAPTER_REPO", "rigidhat/llama-3.2-3b-construction-codecite-v2")
+V3_MODEL_ID = os.environ.get(
+    "V3_MODEL_ID", "rigidhat/llama-4-scout-17b-construction-codecite-v3"
+)
+TOGETHER_API_KEY = os.environ.get("TOGETHER_API_KEY", "")
+FALLBACK_BASE = os.environ.get("FALLBACK_BASE", "meta-llama/Llama-3.2-3B-Instruct")
+FALLBACK_ADAPTER = os.environ.get(
+    "FALLBACK_ADAPTER", "rigidhat/llama-3.2-3b-construction-codecite-v2"
+)
 DATASET_REPO = os.environ.get("DATASET_REPO", "rigidhat/construction-code-corpus-v1")
 MAX_NEW_TOKENS = 384
 RAG_K = 5
@@ -153,59 +165,64 @@ def parse_json(raw: str) -> dict:
     return {}
 
 
-_PIPELINE = None
+_STATE = {"search": None, "verify": None, "fallback": None}
 
 
-def get_pipeline():
-    global _PIPELINE
-    if _PIPELINE is not None:
-        return _PIPELINE
+def get_search_verify():
+    if _STATE["search"] is not None:
+        return _STATE["search"], _STATE["verify"]
     corpus = ensure_corpus()
-    verify = build_verifier(corpus)
-    search = build_bm25(corpus)
+    _STATE["search"] = build_bm25(corpus)
+    _STATE["verify"] = build_verifier(corpus)
+    return _STATE["search"], _STATE["verify"]
+
 
-    print(f"Loading base model: {BASE_MODEL}")
+def get_fallback():
+    """Lazy-load v2 in-Space adapter as the fallback path."""
+    if _STATE["fallback"] is not None:
+        return _STATE["fallback"]
     import torch
     from peft import PeftModel
     from transformers import AutoModelForCausalLM, AutoTokenizer
 
-    tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, trust_remote_code=True)
+    print(f"Loading fallback base: {FALLBACK_BASE}")
+    tokenizer = AutoTokenizer.from_pretrained(FALLBACK_BASE, trust_remote_code=True)
     base = AutoModelForCausalLM.from_pretrained(
-        BASE_MODEL, torch_dtype=torch.float32, trust_remote_code=True
+        FALLBACK_BASE, torch_dtype=torch.float32, trust_remote_code=True
     )
-    print(f"Loading LoRA adapter: {ADAPTER_REPO}")
-    model = PeftModel.from_pretrained(base, ADAPTER_REPO)
-
-    _PIPELINE = {
-        "tokenizer": tokenizer,
-        "model": model,
-        "search": search,
-        "verify": verify,
-        "torch": torch,
-    }
-    return _PIPELINE
-
+    print(f"Loading fallback adapter: {FALLBACK_ADAPTER}")
+    model = PeftModel.from_pretrained(base, FALLBACK_ADAPTER)
+    _STATE["fallback"] = {"tokenizer": tokenizer, "model": model, "torch": torch}
+    return _STATE["fallback"]
 
-def format_candidates(hits) -> str:
-    if not hits:
-        return "(no high-confidence candidates)"
-    return "\n".join(f"- {h['section']}: {h['heading'][:80]}" for h in hits)
 
+def generate_together(prompt: str) -> tuple[str, str]:
+    """Call Together AI hosted endpoint. Returns (text, path_label)."""
+    if not TOGETHER_API_KEY:
+        raise RuntimeError("TOGETHER_API_KEY not set")
+    try:
+        from together import Together
+    except ImportError as e:
+        raise RuntimeError(f"together package not installed: {e}")
+
+    client = Together(api_key=TOGETHER_API_KEY)
+    response = client.chat.completions.create(
+        model=V3_MODEL_ID,
+        messages=[{"role": "user", "content": prompt}],
+        max_tokens=MAX_NEW_TOKENS,
+        temperature=0.0,
+    )
+    return response.choices[0].message.content, "v3 · Llama 4 Scout 17B-16E (Together)"
 
-def predict(narrative: str):
-    if not (narrative or "").strip():
-        return "{}", "(paste an incident narrative first)", "—"
-    t0 = time.time()
-    pipe = get_pipeline()
-    hits = pipe["search"](narrative, k=RAG_K)
-    prompt = PROMPT.format(candidates=format_candidates(hits), narrative=narrative[:1800])
 
+def generate_fallback(prompt: str) -> tuple[str, str]:
+    """Fall back to v2 in-Space."""
+    pipe = get_fallback()
     messages = [{"role": "user", "content": prompt}]
     chat = pipe["tokenizer"].apply_chat_template(
         messages, add_generation_prompt=True, tokenize=False
     )
     inputs = pipe["tokenizer"](chat, return_tensors="pt")
-
     with pipe["torch"].no_grad():
         out = pipe["model"].generate(
             **inputs,
@@ -215,17 +232,41 @@ def predict(narrative: str):
         )
     generated = out[0][inputs["input_ids"].shape[1]:]
     raw = pipe["tokenizer"].decode(generated, skip_special_tokens=True)
-    parsed = parse_json(raw)
+    return raw, "v2 · Llama 3.2 3B (in-Space fallback)"
 
+
+def format_candidates(hits) -> str:
+    if not hits:
+        return "(no high-confidence candidates)"
+    return "\n".join(f"- {h['section']}: {h['heading'][:80]}" for h in hits)
+
+
+def predict(narrative: str):
+    if not (narrative or "").strip():
+        return "{}", "(paste an incident narrative first)", "—", "—"
+    t0 = time.time()
+    search, verify = get_search_verify()
+    hits = search(narrative, k=RAG_K)
+    prompt = PROMPT.format(candidates=format_candidates(hits), narrative=narrative[:1800])
+
+    path_label = ""
+    raw = ""
+    try:
+        raw, path_label = generate_together(prompt)
+    except Exception as e:
+        print(f"Together path failed: {e}. Falling back to v2 in-Space.")
+        raw, path_label = generate_fallback(prompt)
+
+    parsed = parse_json(raw)
     if parsed and "citations" in parsed:
         for c in parsed["citations"]:
-            is_valid, heading = pipe["verify"](c.get("standard", ""))
+            is_valid, heading = verify(c.get("standard", ""))
             c["verified"] = is_valid
             if heading and not c.get("section_heading"):
                 c["section_heading"] = heading
 
     rag_view = "\n".join(f"- {h['section']}: {h['heading'][:80]}" for h in hits)
-    return json.dumps(parsed, indent=2), rag_view, f"{time.time() - t0:.2f}s"
+    return json.dumps(parsed, indent=2), rag_view, f"{time.time() - t0:.2f}s", path_label
 
 
 EXAMPLES = [
@@ -236,14 +277,14 @@ EXAMPLES = [
 
 
 with gr.Blocks(title="Construction Code-Citation") as demo:
-    gr.Markdown("# Construction Code-Citation Model (v2 · AutoScientist)")
+    gr.Markdown("# Construction Code-Citation Model (v3 · Llama 4 Scout 17B-16E · AutoScientist)")
     gr.Markdown(
-        "Llama 3.2 3B fine-tuned by **AutoScientist** on OSHA Severe Injury Reports "
-        "for the [Adaption Labs AutoScientist Challenge](https://adaptionlabs.ai/auto-scientist) "
-        "\"All Other Domains\" category. **77% win rate** vs base Llama 3.2 3B on our test set. "
-        "Given a construction-site incident narrative, returns strict JSON with OIICS hazard "
-        "codes plus OSHA 29 CFR 1926 citations, verifier-grounded against the corpus. "
-        "First request downloads the base model (~6 GB, one-time)."
+        "Llama 4 Scout 17B-16E (MoE) fine-tuned by **AutoScientist** on OSHA Severe "
+        "Injury Reports for the [Adaption Labs AutoScientist Challenge](https://adaptionlabs.ai/auto-scientist) "
+        "\"All Other Domains\" category. Given a construction-site incident narrative, "
+        "returns strict JSON with OIICS hazard codes plus OSHA 29 CFR 1926 citations, "
+        "verifier-grounded against the corpus. **Inference via Together AI hosted "
+        "endpoint** — v2 (Llama 3.2 3B) auto-falls-back if the endpoint is unavailable."
     )
     with gr.Row():
         with gr.Column():
@@ -257,13 +298,16 @@ with gr.Blocks(title="Construction Code-Citation") as demo:
         with gr.Column():
             output_json = gr.Code(label="Hazards + Citations (JSON)", language="json")
             rag_view = gr.Textbox(label="OSHA 1926 RAG candidates (BM25)", lines=6)
-            elapsed = gr.Textbox(label="Latency", lines=1)
+            with gr.Row():
+                elapsed = gr.Textbox(label="Latency", lines=1)
+                path = gr.Textbox(label="Model path", lines=1)
 
-    submit.click(predict, inputs=[narrative], outputs=[output_json, rag_view, elapsed])
+    submit.click(predict, inputs=[narrative], outputs=[output_json, rag_view, elapsed, path])
 
     gr.Markdown(
         "**Artifacts:** "
         "[Dataset](https://huggingface.co/datasets/rigidhat/construction-code-corpus-v1) · "
+        "[v3 Model (Llama 4 Scout 17B-16E · AutoScientist)](https://huggingface.co/rigidhat/llama-4-scout-17b-construction-codecite-v3) · "
         "[v2 Model (Llama 3.2 3B · AutoScientist)](https://huggingface.co/rigidhat/llama-3.2-3b-construction-codecite-v2) · "
         "[v1 Baseline (Qwen 2.5 1.5B)](https://huggingface.co/rigidhat/qwen-2.5-construction-codecite-v1) · "
         "[Source](https://github.com/snakezilla/construction-code-llm)"