Update app.py

app.py

@@ -4,7 +4,8 @@ import torch
 import gradio as gr
 import torch.nn.functional as F
 
-from transformers import GPT2LMHeadModel, CpmTokenizer
+from transformers import GPT2LMHeadModel, BertTokenizer
+
 
 def top_k_top_p_filtering( logits, top_k=0, top_p=0.0, filter_value=-float('Inf') ):
     assert logits.dim() == 1
@@ -52,13 +53,15 @@ def generate(title, context, max_len):
     return result
 
 if __name__ == '__main__':
-    tokenizer = CpmTokenizer(vocab_file="chinese_vocab.model")
-    eod_id = tokenizer.convert_tokens_to_ids("<eod>")
-    sep_id = tokenizer.sep_token_id
-    unk_id = tokenizer.unk_token_id
+    # tokenizer = BertTokenizer(vocab_file="chinese_vocab.model")
+    # eod_id = tokenizer.convert_tokens_to_ids("<eod>")
+    # sep_id = tokenizer.sep_token_id
+    # unk_id = tokenizer.unk_token_id
 
+    tokenizer = BertTokenizer.from_pretrained("supermy/poetry")
     model = GPT2LMHeadModel.from_pretrained("supermy/poetry")
     model.eval()
+    
     gr.Interface(
         fn=generate,
         inputs=[