Add template of prefix, update prompt

app.py

@@ -2,18 +2,15 @@ import os
 import torch
 import gradio as gr
 import spaces
+import json
 from threading import Thread
 from transformers import AutoTokenizer, AutoModelForCausalLM, FineGrainedFP8Config, TextIteratorStreamer
 
 # >>>> CHANGE THIS <<<<
 MODEL_ID = os.getenv("MODEL_ID", "theostos/LLM4Docq-annotator")
-
+RESULT_JSON_PATH = os.getenv("RESULT_JSON_PATH", "result.json")
 # Matches your training style: messages=[{"role":"user","content": template.format(term=..., dependencies=...)}]
-INSTRUCTION_TEMPLATE = (
-    "You are a Rocq code annotator. Given the Coq term and its dependencies, "
-    "produce helpful inline comments and explanations.\n\n"
-    "Term:\n{term}\n\nDependencies:\n{dependencies}\n"
-)
+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}"
 
 HF_TOKEN = os.getenv("HF_TOKEN")
 
@@ -37,17 +34,32 @@ def load_model():
         )
     return _model
 
-def build_messages(term: str, deps: str):
-    content = INSTRUCTION_TEMPLATE.format(term=term, dependencies=deps)
+def build_messages(prefix: str, source: str):
+    content = INSTRUCTION_TEMPLATE.format(prefix=prefix, source=source)
     return [{"role": "user", "content": content}]
 
+def load_prefixes(path=RESULT_JSON_PATH):
+    try:
+        with open(path, "r", encoding="utf-8") as f:
+            data = json.load(f)
+        if not isinstance(data, dict):
+            raise ValueError("result.json must be a JSON object mapping keys -> prefix strings.")
+        # coerce to str->str
+        return {str(k): str(v) for k, v in data.items()}
+    except Exception as e:
+        print(f"[warn] Could not load {path}: {e}")
+        return {}
+
+PREFIXES = load_prefixes()
+PREFIX_KEYS = sorted(PREFIXES.keys())
+
 # Estimate duration for ZeroGPU (default is 60s). Shorter = better queue priority.
-def _duration(term, deps, temperature, top_p, max_new_tokens, repetition_penalty):
+def _duration(term, deps, temperature, top_p, max_new_tokens):
     # crude: ~2.5 tok/s + 30s headroom
     return int(min(300, max(60, (int(max_new_tokens) / 2.5) + 30)))
 
 @spaces.GPU(duration=_duration)
-def generate(term, deps, temperature, top_p, max_new_tokens, repetition_penalty):
+def generate(term, deps, temperature, top_p, max_new_tokens):
     model = load_model()
     device = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -79,23 +91,68 @@ def generate(term, deps, temperature, top_p, max_new_tokens, repetition_penalty)
         out += token
         yield f"```rocq\n{out}\n```"
 
-with gr.Blocks(title="Rocq Annotator (ZeroGPU)") as demo:
-    gr.Markdown("# Rocq annotator\nThe model will produce annotated Rocq code.")
+def set_prefix_from_key(key: str) -> str:
+    return PREFIXES.get(key, "") if key else ""
+
+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**."
+    )
+
     with gr.Row():
-        term = gr.Textbox(label="Prefix", lines=100, placeholder="Paste the prefix to use")
-        deps = gr.Textbox(label="To annotate", lines=8, placeholder="The code to annotate")
+        dropdown = gr.Dropdown(
+            choices=PREFIX_KEYS,
+            label="Choose a prefix example (from result.json)",
+            allow_custom_value=False,
+            value=None,
+        )
+
+        reload_btn = gr.Button("Reload result.json", variant="secondary")
+
     with gr.Row():
-        temperature = gr.Slider(0.0, 1.5, value=0.7, step=0.05, label="Temperature")
+        prefix_box = gr.Code(
+            label="Prefix (context; auto-filled from dropdown, then editable)",
+            language="coq",
+            interactive=True,
+            lines=18,
+        )
+        target_box = gr.Code(
+            label="Target snippet (contains [PREDICT_STATEMENT] / [PREDICT_DOCSTRING])",
+            language="coq",
+            interactive=True,
+            lines=18,
+            placeholder="Paste the code to annotate…",
+        )
+
+    with gr.Row():
+        temperature = gr.Slider(0.0, 1.5, value=0.2, 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(256, 8192, value=4096, step=32, label="max_new_tokens")
+        max_new = gr.Slider(32, 512, value=128, step=32, label="max_new_tokens")
+
     out = gr.Markdown(label="Annotated Rocq")
-    btn = gr.Button("Annotate")
+    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
+        PREFIXES = load_prefixes()
+        PREFIX_KEYS = sorted(PREFIXES.keys())
+        # return updated dropdown (choices) and a notice
+        return gr.update(choices=PREFIX_KEYS), gr.update(value="Reloaded result.json.")
+    notice = gr.Markdown("")
+    reload_btn.click(_reload, inputs=None, outputs=[dropdown, notice])
+
     btn.click(
         generate,
-        inputs=[term, deps, temperature, top_p, max_new],
+        inputs=[prefix_box, target_box, temperature, top_p, max_new],
         outputs=out,
-        concurrency_limit=1,  # cooperate with ZeroGPU queues
+        concurrency_limit=1,
     )
+
     demo.queue(max_size=20, default_concurrency_limit=1)
 
 if __name__ == "__main__":