25 lines
1.3 KiB
Python
25 lines
1.3 KiB
Python
|
#!/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()
|