File size: 1,270 Bytes
c83c617 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 | """Read-only V4 next-piece, beam and optional candidate scoring diagnostics."""
import argparse
import json
import sys
from pathlib import Path
ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "src"))
from vimeml.training.infer_v4 import JapaneseLMV4, score_candidates_cached
from vimeml.training.evaluate_ime import beam_suggestions
if __name__ == "__main__":
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, required=True)
parser.add_argument("--tokenizer", type=Path, default=ROOT / "artifacts/tokenizers/v4-a/ja-unigram-8192")
parser.add_argument("--prompt", default="今日は")
parser.add_argument("--candidate", action="append", default=[])
parser.add_argument("--device", choices=("cpu", "cuda"), default="cpu")
args = parser.parse_args()
lm = JapaneseLMV4(args.checkpoint, args.tokenizer, args.device)
result = {"metadata": lm.metadata, "next_pieces": lm.next_pieces(args.prompt),
"beam": beam_suggestions(lm, args.prompt, max_tokens=3, width=3, count=2)}
if args.candidate:
result["candidate_scores"] = score_candidates_cached(lm, args.prompt, args.candidate)
print(json.dumps(result, ensure_ascii=False, indent=2))
|