Create new file

correct.py

@@ -0,0 +1,39 @@
+# -*- coding: utf-8 -*-
+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()