infer.py 2.8 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768
  1. import torch
  2. from transformers import PreTrainedTokenizerFast
  3. from transformers.models.bart import BartForConditionalGeneration
  4. import os
  5. root_path = os.getcwd()
  6. summary_general = 'summary/kobart_summary'
  7. summary_estate = 'summary/kobart_summary_gen'
  8. model_estate = BartForConditionalGeneration.from_pretrained(os.path.join(root_path, summary_estate))
  9. model_gen = BartForConditionalGeneration.from_pretrained(os.path.join(root_path, summary_general))
  10. #tokenizer = PreTrainedTokenizerFast.from_pretrained('gogamza/kobart-base-v1')
  11. # tokenizer = PreTrainedTokenizerFast(tokenizer_file = 'tokenizer.json')
  12. tokenizer = PreTrainedTokenizerFast(tokenizer_file = os.path.join(root_path, 'summary/tokenizer.json'))
  13. # text = input()
  14. import pandas as pd
  15. def infer():
  16. test = pd.read_csv("test_label.tsv", sep = '\t', encoding = "utf-8")
  17. for text in test["text"]:
  18. if text:
  19. input_ids = tokenizer.encode(text)
  20. input_ids = torch.tensor(input_ids)
  21. input_ids = input_ids.unsqueeze(0)
  22. output = model.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
  23. output = tokenizer.decode(output[0], skip_special_tokens=True)
  24. print(">>>"+"원문 ---- "+text)
  25. print("\n"+ "######################################################################" + "\n")
  26. print(">>>"+"요약 ---- "+output)
  27. print("\n"+ "######################################################################" + "\n")
  28. def summary(text):
  29. input_ids = tokenizer.encode(text)
  30. input_ids = torch.tensor(input_ids)
  31. input_ids = input_ids.unsqueeze(0)
  32. output = model_estate.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
  33. output = tokenizer.decode(output[0], skip_special_tokens=True)
  34. return output
  35. def summary_with_model(text, idx):
  36. input_ids = tokenizer.encode(text)
  37. input_ids = torch.tensor(input_ids)
  38. input_ids = input_ids.unsqueeze(0)
  39. if idx == 0:
  40. output = model_estate.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
  41. else:
  42. output = model_gen.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
  43. output = tokenizer.decode(output[0], skip_special_tokens=True)
  44. return output
  45. if __name__ == '__main__':
  46. print(summary('사장님 분위기 좀 알려주세요. 요즘에 매매 전세 전반적인 분위기가 어때요'\
  47. '그래도 이제 수능 끝나고 좀 나아져야 되는데 매물은 없고 지금 현재 손님도 안 찾고 그래요'\
  48. '물건을 좀 찾기는 하는데 물건이 안 나온단 얘기죠'\
  49. '전세도 그런가요'\
  50. '전세도그래요 '\
  51. '전세나 매물이 많이 나오지 않으세요.'\
  52. '예 안 나옵니다. '\
  53. '찾는 사람은 간혹 있어도 전세가가 약간 올랐어요.'\
  54. '알겠습니다.' \
  55. '사장님 시세는 그대로 하겠습니다. 감사합니다.'))