Update correct.py

correct.py

@@ -12,27 +12,27 @@ 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():
+    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:
-                q = input('원래문장: ').strip()
-                if q == 'quit':
+                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 = ''
-                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()}")
+                a += gen.replace('▁', ' ')
+            print(f"교정: {a.strip()}")
 
 
 if __name__ == "__main__":