| """FP32 candidate scoring and phrase prediction using the included VimeML 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("--context", default="我们明天一起去") |
| parser.add_argument("--candidate", action="append", help="Candidate string; repeat to rerank.") |
| parser.add_argument("--previous", action="append", default=[], |
| help="Earlier turn of the conversation, oldest first; repeat for several.") |
| parser.add_argument("--threads", type=int, default=4) |
| args = parser.parse_args() |
| if args.threads < 1: |
| parser.error("threads must be positive.") |
| sys.path.insert(0, str(ROOT / "source/src")) |
| import torch |
| from vimeml.benchmarks import predict_zh |
| from vimeml.benchmarks.zh13_lm_scores import score_item |
| from vimeml.deployment.zh_v1 import BundleLMZh |
|
|
| torch.set_num_threads(args.threads) |
| bundle = BundleLMZh(ROOT / "inference") |
| if args.candidate: |
| result = score_item(bundle.model, bundle.processor, args.context, args.candidate) |
| result["ranked_candidates"] = [ |
| {"candidate": c, "log_probability_sum": s} |
| for s, c in sorted(zip(result["scores"], args.candidate), key=lambda x: -x[0]) |
| ] |
| else: |
| lm = predict_zh.VimeLM(bundle.model, bundle.processor, "cpu", bundle.special) |
| predictions = predict_zh.predict(lm, args.previous, args.context) |
| result = {"context": args.context, "previous_turns": args.previous, |
| "predictions": [{"text": t, "log_probability_sum": s} for t, s in predictions]} |
| print(json.dumps(result, ensure_ascii=False, indent=2)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
