| """FP32 inference and candidate scoring using the included VimeML V2 source.""" |
|
|
| import argparse |
| import json |
| import sys |
| from pathlib import Path |
|
|
| ROOT = Path(__file__).resolve().parent |
|
|
|
|
| def main(): |
| parser = argparse.ArgumentParser(description=__doc__) |
| parser.add_argument("--prompt", default="今日は雨が降っているので、") |
| parser.add_argument("--candidate", action="append", help="Candidate string; repeat to rerank.") |
| parser.add_argument("--max-new-tokens", type=int, default=8) |
| parser.add_argument("--threads", type=int, default=4) |
| args = parser.parse_args() |
| if args.threads < 1 or args.max_new_tokens < 1: |
| parser.error("threads and max-new-tokens must be positive.") |
| sys.path.insert(0, str(ROOT / "source/src")) |
| import torch |
| from vimeml.deployment.v2 import BundleLM |
| from vimeml.training.evaluate_ime import score_candidates |
|
|
| torch.set_num_threads(args.threads) |
| lm = BundleLM(ROOT / "inference") |
| if args.candidate: |
| result = score_candidates(lm, args.prompt, args.candidate) |
| result["ranked_candidates"] = sorted( |
| result["candidates"], key=lambda item: item["log_probability_sum"], reverse=True |
| ) |
| else: |
| result = lm.generate(args.prompt, max_new_tokens=args.max_new_tokens, temperature=0) |
| print(json.dumps(result, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
