1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768 |
- import torch
- from transformers import PreTrainedTokenizerFast
- from transformers.models.bart import BartForConditionalGeneration
- import os
- root_path = os.getcwd()
- summary_general = 'summary/kobart_summary'
- summary_estate = 'summary/kobart_summary_gen'
- model_estate = BartForConditionalGeneration.from_pretrained(os.path.join(root_path, summary_estate))
- model_gen = BartForConditionalGeneration.from_pretrained(os.path.join(root_path, summary_general))
- #tokenizer = PreTrainedTokenizerFast.from_pretrained('gogamza/kobart-base-v1')
- # tokenizer = PreTrainedTokenizerFast(tokenizer_file = 'tokenizer.json')
- tokenizer = PreTrainedTokenizerFast(tokenizer_file = os.path.join(root_path, 'summary/tokenizer.json'))
- # text = input()
- import pandas as pd
- def infer():
- test = pd.read_csv("test_label.tsv", sep = '\t', encoding = "utf-8")
- for text in test["text"]:
- if text:
- input_ids = tokenizer.encode(text)
- input_ids = torch.tensor(input_ids)
- input_ids = input_ids.unsqueeze(0)
- output = model.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
- output = tokenizer.decode(output[0], skip_special_tokens=True)
- print(">>>"+"원문 ---- "+text)
- print("\n"+ "######################################################################" + "\n")
- print(">>>"+"요약 ---- "+output)
- print("\n"+ "######################################################################" + "\n")
- def summary(text):
- input_ids = tokenizer.encode(text)
- input_ids = torch.tensor(input_ids)
- input_ids = input_ids.unsqueeze(0)
- output = model_estate.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
- output = tokenizer.decode(output[0], skip_special_tokens=True)
- return output
- def summary_with_model(text, idx):
- input_ids = tokenizer.encode(text)
- input_ids = torch.tensor(input_ids)
- input_ids = input_ids.unsqueeze(0)
- if idx == 0:
- output = model_estate.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
- else:
- output = model_gen.generate(input_ids, eos_token_id=1, max_length=512, num_beams=7)
- output = tokenizer.decode(output[0], skip_special_tokens=True)
- return output
- if __name__ == '__main__':
- print(summary('사장님 분위기 좀 알려주세요. 요즘에 매매 전세 전반적인 분위기가 어때요'\
- '그래도 이제 수능 끝나고 좀 나아져야 되는데 매물은 없고 지금 현재 손님도 안 찾고 그래요'\
- '물건을 좀 찾기는 하는데 물건이 안 나온단 얘기죠'\
- '전세도 그런가요'\
- '전세도그래요 '\
- '전세나 매물이 많이 나오지 않으세요.'\
- '예 안 나옵니다. '\
- '찾는 사람은 간혹 있어도 전세가가 약간 올랐어요.'\
- '알겠습니다.' \
- '사장님 시세는 그대로 하겠습니다. 감사합니다.'))
|