初始化项目,由ModelHub XC社区提供模型
Model: Luimas/claim-extractor-detective-qwen3b Source: Original Platform
This commit is contained in:
24
scripts/inference.py
Normal file
24
scripts/inference.py
Normal file
@@ -0,0 +1,24 @@
|
||||
#!/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()
|
||||
Reference in New Issue
Block a user