Add draw from test set, pull fp8 model

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+result.json filter=lfs diff=lfs merge=lfs -text
+test.json filter=lfs diff=lfs merge=lfs -text

app.py

@@ -3,12 +3,15 @@ import torch
 import gradio as gr
 import spaces
 import json
+import random  # <<< NEW
 from threading import Thread
-from transformers import AutoTokenizer, AutoModelForCausalLM, FineGrainedFP8Config, TextIteratorStreamer
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer
 
 # >>>> CHANGE THIS <<<<
-MODEL_ID = os.getenv("MODEL_ID", "theostos/LLM4Docq-annotator")
+MODEL_ID = os.getenv("MODEL_ID", "theostos/LLM4Docq-annotator-fp8")
 RESULT_JSON_PATH = os.getenv("RESULT_JSON_PATH", "result.json")
+# NEW: path to your test set (list[dict])
+TEST_JSON_PATH = os.getenv("TEST_JSON_PATH", "test.json")
 # Matches your training style: messages=[{"role":"user","content": template.format(term=..., dependencies=...)}]
 INSTRUCTION_TEMPLATE = "You are given a Coq source file along with an optional prefix.\n\n- The **prefix** contains lines that appear *before* the current chunk of code. It provides contextual information to help you understand the surrounding definitions, imports, and notation.\n- The **source** contains the chunk of code you must annotate and complete.\n\nSome parts of the code contain special placeholders:\n\n- [PREDICT_DOCSTRING]: This placeholder appears before an element. You must replace it with a descriptive comment (in Coq comment syntax (* ... *)) that explains what the element does.\n\n- [PREDICT_STATEMENT]: This placeholder appears after an explanatory comment. You must replace it with a valid Coq statement or definition that matches the meaning of the preceding comment.\n\nYour task is to rewrite the entire Coq source chunk, replacing all placeholders with appropriate content, while preserving all other parts of the source code exactly as they are.\n\n### Guidelines\n1. The **prefix** is only provided for context — do **not** modify it or include it in your output.\n2. Rewrite only the **source** content.\n3. Keep all existing Coq syntax, imports, and formatting intact.\n4. Replace [PREDICT_DOCSTRING] with a natural-language description of the next element.\n5. Replace [PREDICT_STATEMENT] with a complete and syntactically correct Coq statement (definition, lemma, theorem, etc.) that corresponds to the immediately preceding comment.\n6. Ensure the generated statements are consistent with the style and logic suggested by the prefix and surrounding code.\n7. Do not add or remove any lines except to substitute the placeholders.\n\n### Output format\nReturn **only** the full rewritten Coq source chunk (without the prefix), with all placeholders replaced.\n\nHere is the context and source:\n\n## Prefix:\n{prefix}\n\n## Source:\n{source}"
 
@@ -18,8 +21,6 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, use_fast=Tru
 if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None:
     tokenizer.pad_token_id = tokenizer.eos_token_id
 
-quant_config = FineGrainedFP8Config()
-
 _model = None
 def load_model():
     global _model
@@ -28,9 +29,8 @@ def load_model():
             MODEL_ID,
             token=HF_TOKEN,
             device_map="auto",
-            dtype="auto",                 # load base weights in their optimal dtype
-            quantization_config=quant_config,  # <-- FP8 quantization
-            trust_remote_code=True,
+            dtype="auto",
+            trust_remote_code=True
         )
     return _model
 
@@ -50,9 +50,39 @@ def load_prefixes(path=RESULT_JSON_PATH):
         print(f"[warn] Could not load {path}: {e}")
         return {}
 
+# --- NEW: test set loader + helpers ---
+def load_test_examples(path=TEST_JSON_PATH):
+    """
+    Expects a JSON list of dicts with keys:
+      - 'prefix'
+      - 'partially_annotated_last_target'
+      - 'fully_annotated_last_target'
+    Returns a cleaned list with {'prefix','target','truth'} strings.
+    """
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, list):
+            raise ValueError("Test set JSON must be a list of objects.")
+        cleaned = []
+        for i, ex in enumerate(data):
+            if not isinstance(ex, dict):
+                continue
+            prefix = str(ex.get("prefix", ""))
+            target = str(ex.get("partially_annotated_last_target", ""))
+            truth  = str(ex.get("fully_annotated_last_target", ""))
+            cleaned.append({"prefix": prefix, "target": target, "truth": truth})
+        print(f"[info] Loaded {len(cleaned)} test examples from {path}")
+        return cleaned
+    except Exception as e:
+        print(f"[warn] Could not load test set {path}: {e}")
+        return []
+
 PREFIXES = load_prefixes()
 PREFIX_KEYS = sorted(PREFIXES.keys())
 
+TEST_EXAMPLES = load_test_examples()  # <<< NEW
+
 # Estimate duration for ZeroGPU (default is 60s). Shorter = better queue priority.
 def _duration(term, deps, temperature, top_p, max_new_tokens):
     # crude: ~2.5 tok/s + 30s headroom
@@ -94,11 +124,27 @@ def generate(term, deps, temperature, top_p, max_new_tokens):
 def set_prefix_from_key(key: str) -> str:
     return PREFIXES.get(key, "") if key else ""
 
+# NEW: sample a random test example
+def _sample_test_example():
+    if not TEST_EXAMPLES:
+        # Return empty prefix/target and a notice in the truth box
+        return "", "", "No test examples loaded. Set TEST_JSON_PATH or add test.json at repo root."
+    ex = random.choice(TEST_EXAMPLES)
+    truth_md = f"```rocq\n{ex['truth']}\n```" if ex["truth"] else ""
+    return ex["prefix"], ex["target"], truth_md
+
+# NEW: hot-reload the test set
+def _reload_test_set():
+    global TEST_EXAMPLES
+    TEST_EXAMPLES = load_test_examples()
+    return gr.update(value=f"Reloaded {len(TEST_EXAMPLES)} test examples from {TEST_JSON_PATH}.")
+
 with gr.Blocks(title="Rocq Annotator (ZeroGPU, FP8)") as demo:
     gr.Markdown(
         "# Rocq annotator\n"
         "Pick a **prefix** example from the dropdown to auto-fill the Prefix editor, "
-        "then write a **target snippet** (with [PREDICT_STATEMENT]/[PREDICT_DOCSTRING] tags) and click **Annotate**."
+        "then write a **target snippet** (with [PREDICT_STATEMENT]/[PREDICT_DOCSTRING] tags) and click **Annotate**.\n\n"
+        "You can also use **🎲 Draw test example** to pull a sample from the test set; the **Baseline (truth)** panel shows the expected annotated result."
     )
 
     with gr.Row():
@@ -110,6 +156,8 @@ with gr.Blocks(title="Rocq Annotator (ZeroGPU, FP8)") as demo:
         )
 
         reload_btn = gr.Button("Reload result.json", variant="secondary")
+        sample_btn = gr.Button("🎲 Draw test example", variant="secondary")
+        reload_test_btn = gr.Button("Reload test set", variant="secondary")
 
     with gr.Row():
         prefix_box = gr.Code(
@@ -126,15 +174,20 @@ with gr.Blocks(title="Rocq Annotator (ZeroGPU, FP8)") as demo:
         )
 
     with gr.Row():
-        temperature = gr.Slider(0.0, 1.5, value=0.2, step=0.05, label="Temperature")
+        temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
         top_p = gr.Slider(0.1, 1.0, value=0.9, step=0.05, label="top_p")
         max_new = gr.Slider(32, 512, value=128, step=32, label="max_new_tokens")
 
-    out = gr.Markdown(label="Annotated Rocq")
+    # Output panels: model vs baseline/truth
+    with gr.Row():
+        out = gr.Markdown(label="Annotated Rocq")
+        truth_md = gr.Markdown(label="Baseline (truth)")  # <<< NEW
+
     btn = gr.Button("Annotate", variant="primary")
 
     # --- wiring ---
     dropdown.change(set_prefix_from_key, inputs=dropdown, outputs=prefix_box)
+
     # Optional: hot reload result.json without restarting Space
     def _reload():
         global PREFIXES, PREFIX_KEYS
@@ -145,6 +198,11 @@ with gr.Blocks(title="Rocq Annotator (ZeroGPU, FP8)") as demo:
     notice = gr.Markdown("")
     reload_btn.click(_reload, inputs=None, outputs=[dropdown, notice])
 
+    # NEW: wire test set actions
+    test_notice = gr.Markdown("")
+    sample_btn.click(_sample_test_example, inputs=None, outputs=[prefix_box, target_box, truth_md])
+    reload_test_btn.click(_reload_test_set, inputs=None, outputs=test_notice)
+
     btn.click(
         generate,
         inputs=[prefix_box, target_box, temperature, top_p, max_new],

test.json

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:547fd7e37c43ef05a1681a60192234a31ebb95efad32e829937ebf77ae5140b3
+size 16240662