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('사장님 분위기 좀 알려주세요. 요즘에 매매 전세 전반적인 분위기가 어때요'\ '그래도 이제 수능 끝나고 좀 나아져야 되는데 매물은 없고 지금 현재 손님도 안 찾고 그래요'\ '물건을 좀 찾기는 하는데 물건이 안 나온단 얘기죠'\ '전세도 그런가요'\ '전세도그래요 '\ '전세나 매물이 많이 나오지 않으세요.'\ '예 안 나옵니다. '\ '찾는 사람은 간혹 있어도 전세가가 약간 올랐어요.'\ '알겠습니다.' \ '사장님 시세는 그대로 하겠습니다. 감사합니다.'))