| |
| import torch |
| from transformers import AutoTokenizer, GPT2LMHeadModel |
|
|
| O_TKN = '<origin>' |
| C_TKN = '<correct>' |
| BOS = "</s>" |
| EOS = "</s>" |
| PAD = "<pad>" |
| MASK = '<unused0>' |
| SENT = '<unused1>' |
|
|
|
|
| def chat(): |
| tokenizer = AutoTokenizer.from_pretrained('skt/kogpt2-base-v2', |
| eos_token=EOS, unk_token='<unk>', |
| pad_token=PAD, mask_token=MASK) |
| model = GPT2LMHeadModel.from_pretrained('Moo/kogpt2-proofreader') |
| with torch.no_grad(): |
| while True: |
| q = input('원래문장: ').strip() |
| if q == 'quit': |
| break |
| a = '' |
| while True: |
| input_ids = torch.LongTensor(tokenizer.encode(O_TKN + q + C_TKN + a)).unsqueeze(dim=0) |
| pred = model(input_ids) |
| gen = tokenizer.convert_ids_to_tokens( |
| torch.argmax( |
| pred[0], |
| dim=-1).squeeze().numpy().tolist())[-1] |
| if gen == EOS: |
| break |
| a += gen.replace('▁', ' ') |
| print(f"교정: {a.strip()}") |
|
|
|
|
| if __name__ == "__main__": |
| chat() |
|
|
