JJ-1-2B / example.py
LewisLuo's picture
Add model weights and inference
dbf752e verified
Raw History Blame Contribute Delete
1.76 kB
"""Run the same Japanese request through Transformers or llama.cpp."""
import argparse
import json
from pathlib import Path
STATE = "洗面所の天井から水が落ち続けていて、床が水浸しです。至急見に来てください。"
QUESTIONS = {
"category": {
"type": "choice",
"instructions": "問い合わせの種類は?",
"criteria": {"water": "水漏れ", "noise": "騒音", "other": "その他"},
},
"needs_visit": {"type": "noul", "instructions": "現地での対応が必要ですか?"},
"urgency": {
"type": "score",
"instructions": "対応の緊急度は?",
"criteria": ["急がない", "数日以内", "今すぐ"],
},
}
def main():
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--backend", choices=["transformers", "gguf"], default="transformers")
parser.add_argument("--model", default=str(Path(__file__).resolve().parent))
parser.add_argument("--device", choices=["cpu", "cuda"], default="cpu")
parser.add_argument("--url", default="http://127.0.0.1:8081")
args = parser.parse_args()
if args.backend == "gguf":
from jj1 import GGUF
model = GGUF(args.model, url=args.url)
try:
answers, usage = model.decide(STATE, QUESTIONS)
finally:
model.close()
else:
import torch
from jj1 import JJ1
if args.device == "cpu":
torch.set_num_threads(4)
model = JJ1(args.model, images=False, device=args.device, max_length=4096)
answers, usage = model.decide(STATE, QUESTIONS)
print(json.dumps({"answers": answers, "usage": usage}, ensure_ascii=False, indent=2))
if __name__ == "__main__":
main()