import os import torch import gradio as gr import spaces import json import random # <<< NEW from threading import Thread from transformers import AutoTokenizer, AutoModelForCausalLM, TextIteratorStreamer # >>>> CHANGE THIS <<<< 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}" HF_TOKEN = os.getenv("HF_TOKEN") tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, token=HF_TOKEN, use_fast=True) if tokenizer.pad_token_id is None and tokenizer.eos_token_id is not None: tokenizer.pad_token_id = tokenizer.eos_token_id _model = None def load_model(): global _model if _model is None: _model = AutoModelForCausalLM.from_pretrained( MODEL_ID, token=HF_TOKEN, device_map="auto", dtype="auto", trust_remote_code=True ) return _model 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 {} # --- 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 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): model = load_model() device = "cuda" if torch.cuda.is_available() else "cpu" messages = build_messages(term, deps) inputs = tokenizer.apply_chat_template( messages, tokenize=True, add_generation_prompt=True, return_tensors="pt", ).to(device) streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True) gen_kwargs = dict( inputs=inputs, max_new_tokens=int(max_new_tokens), temperature=float(temperature), top_p=float(top_p), do_sample=True, streamer=streamer, pad_token_id=tokenizer.pad_token_id, eos_token_id=tokenizer.eos_token_id, ) thread = Thread(target=model.generate, kwargs=gen_kwargs) thread.start() out = "" for token in streamer: # stream tokens to UI out += token yield f"```rocq\n{out}\n```" 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**.\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(): 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") 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( label="Prefix (context; auto-filled from dropdown, then editable)", language=None, interactive=True, lines=18, ) target_box = gr.Code( label="Target snippet (contains [PREDICT_STATEMENT] / [PREDICT_DOCSTRING])", language=None, interactive=True, lines=18, ) with gr.Row(): 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") # 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 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]) # 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], outputs=out, concurrency_limit=1, ) demo.queue(max_size=20, default_concurrency_limit=1) if __name__ == "__main__": demo.launch()