#!/usr/bin/env python
# Local grammar-constrained inference (always-valid JSON).  pip install llama-cpp-python
import argparse, glob, json, os, sys
from llama_cpp import Llama, LlamaGrammar
HERE = os.path.dirname(os.path.abspath(__file__)); ROOT = os.path.dirname(HERE)
def _find(name):
    for d in (ROOT, HERE):
        hits = glob.glob(os.path.join(d, name))
        if hits: return hits[0]
    raise FileNotFoundError(name)
GGUF = _find("*.gguf")
PROMPT = open(_find("prompt.txt"), encoding="utf-8").read()
GRAMMAR = LlamaGrammar.from_string(open(_find("claim.gbnf"), encoding="utf-8").read())
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--text"); ap.add_argument("--file"); ap.add_argument("--n_gpu_layers", type=int, default=-1)
    a = ap.parse_args()
    text = open(a.file, encoding="utf-8").read() if a.file else (a.text or sys.stdin.read())
    llm = Llama(model_path=GGUF, n_ctx=4096, n_gpu_layers=a.n_gpu_layers, verbose=False)
    out = llm.create_chat_completion(messages=[{"role": "user", "content": PROMPT + text[:4000]}],
                                     grammar=GRAMMAR, temperature=0.0, max_tokens=768)
    print(json.dumps(json.loads(out["choices"][0]["message"]["content"]), indent=2, ensure_ascii=False))
if __name__ == "__main__":
    main()