fix tokenizer

app.py

@@ -1,10 +1,7 @@
 import torch
 import string
 import streamlit as st
-from transformers import GPT2LMHeadModel
-from tokenizers import Tokenizer
-
-
+from transformers import GPT2LMHeadModel, PreTrainedTokenizerFast
 
 
 @st.cache
@@ -13,7 +10,13 @@ def get_model():
     model.eval()
     return model
 
-tokenizer = Tokenizer.from_file('skt/kogpt2-base-v2')
+tokenizer = PreTrainedTokenizerFast.from_pretrained("skt/kogpt2-base-v2",
+                                                    bos_token='</s>',
+                                                    eos_token='</s>',
+                                                    unk_token='<unk>',
+                                                    pad_token='<pad>',
+                                                    mask_token='<mask>')
+
 
 default_text = "현대인들은 왜 항상 불안해 할까?"
 
@@ -39,26 +42,16 @@ st.markdown("""
 
 text = st.text_area("Input Text:", value=default_text)
 st.write(text)
-st.markdown("""
-> *현재 2core 인스턴스에서 예측이 진행되어 다소 느릴 수 있음*
-""")
 punct = ('!', '?', '.')
 
 if text:
     st.markdown("## Predict")
     with st.spinner('processing..'):
         print(f'input > {text}') 
-        input_ids = tokenizer.encode(text).ids
+        input_ids = tokenizer(text)['input_ids']
         gen_ids = model.generate(torch.tensor([input_ids]),
                                     max_length=128,
-                                    repetition_penalty=2.0,
-                                    # num_beams=2,
-                                    # length_penalty=1.0,
-                                    use_cache=True,
-                                    pad_token_id=tokenizer.token_to_id('<pad>'),
-                                    eos_token_id=tokenizer.token_to_id('</s>'),
-                                    bos_token_id=tokenizer.token_to_id('</s>'),
-                                    bad_words_ids=[[tokenizer.token_to_id('<unk>')] ])
+                                    repetition_penalty=2.0)
         generated = tokenizer.decode(gen_ids[0,:].tolist()).strip()
         if generated != '' and generated[-1] not in punct:
             for i in reversed(range(len(generated))):