add zh

app.py

@@ -25,15 +25,19 @@ def download_file(filename):
 
 NLTK = nltk_load(download_file('english.pickle'))
 sent_cut_en = NLTK.tokenize
-LR_GLTR_EN, LR_PPL_EN = [
+LR_GLTR_EN, LR_PPL_EN, LR_GLTR_ZH, LR_PPL_ZH = [
     pickle.load(open(download_file(f'{lang}-gpt2-{name}.pkl'), 'rb'))
-    for lang, name in [('en', 'gltr'), ('en', 'ppl')]
+    for lang, name in [('en', 'gltr'), ('en', 'ppl'), ('zh', 'gltr'), ('zh', 'ppl')]
 ]
 
 NAME_EN = 'gpt2'
 TOKENIZER_EN = GPT2Tokenizer.from_pretrained(NAME_EN)
 MODEL_EN = GPT2LMHeadModel.from_pretrained(NAME_EN)
 
+NAME_ZH = 'IDEA-CCNL/Wenzhong-GPT2-110M'
+TOKENIZER_ZH = GPT2Tokenizer.from_pretrained(NAME_ZH)
+MODEL_ZH = GPT2LMHeadModel.from_pretrained(NAME_ZH)
+
 
 # code borrowed from https://github.com/blmoistawinde/HarvestText
 def sent_cut_zh(para: str) -> List[str]:
@@ -143,7 +147,7 @@ def predict_en(text: str) -> List:
 def predict_zh(text: str) -> List:
     with torch.no_grad():
         feat = gpt2_features(text, TOKENIZER_ZH, MODEL_ZH, sent_cut_zh)
-    out = lr_predict(*feat, None, None, ['人类', 'ChatGPT'])
+    out = lr_predict(*feat, LR_GLTR_ZH, LR_PPL_ZH, ['人类', 'ChatGPT'])
     return out
 
 
@@ -208,10 +212,10 @@ with gr.Blocks() as demo:
             value="对于OpenAI大力出奇迹的工作，自然每个人都有自己的看点。我自己最欣赏的地方是ChatGPT如何解决 “AI校正(Alignment)“这个问题。这个问题也是我们课题组这两年在探索的学术问题之一。"
         )
         button2 = gr.Button("🤖 预测!")
-        gr.Markdown("GLTR")
+        gr.Markdown("GLTR (中文测试集准确率 86.39%)")
         label2_gltr = gr.Textbox(lines=1, label='预测结果 🎃')
         score2_gltr = gr.Textbox(lines=1, label='模型概率')
-        gr.Markdown("PPL")
+        gr.Markdown("PPL (中文测试集准确率 59.04%, 持续优化中...)")
         label2_ppl = gr.Textbox(lines=1, label='PPL 预测结果 🎃')
         score2_ppl = gr.Textbox(lines=1, label='PPL 模型概率')