example.py

"""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()