predict.py 1.4 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546
  1. # -*- coding: utf-8 -*-
  2. """
  3. @author:XuMing(xuming624@qq.com)
  4. @description:
  5. """
  6. import argparse
  7. import sys
  8. sys.path.append('../..')
  9. from pycorrector.deepcontext.deepcontext_corrector import DeepContextCorrector
  10. def main():
  11. parser = argparse.ArgumentParser()
  12. parser.add_argument("--output_dir", default="./outputs-deepcontext-lm", type=str, help="Dir for model save.")
  13. parser.add_argument("--max_length", default=1024, type=int, help="The maximum total input sequence length")
  14. args = parser.parse_args()
  15. print(args)
  16. m = DeepContextCorrector(args.output_dir, max_length=args.max_length)
  17. inputs = [
  18. '老是较书。',
  19. '感谢等五分以后,碰到一位很棒的奴生跟我可聊。',
  20. '遇到一位很棒的奴生跟我聊天。',
  21. '遇到一位很美的女生跟我疗天。',
  22. '他们只能有两个选择:接受降新或自动离职。',
  23. '王天华开心得一直说话。'
  24. ]
  25. for i in inputs:
  26. output = m.correct(i)
  27. print('input :', i)
  28. print('predict:', output)
  29. print()
  30. # result:
  31. # input : 老是较书。
  32. # predict: 老师教书。 [('是', '师', 1, 2), ('较', '教', 2, 3)]
  33. #
  34. # input : 感谢等五分以后,碰到一位很棒的奴生跟我可聊。
  35. # predict: 感谢等五分以后,碰到一位很棒的女生跟我可聊。 [('奴', '女', 15, 16)]
  36. if __name__ == "__main__":
  37. main()