Update app.py

app.py

@@ -1,25 +1,79 @@
 import gradio as gr
-from transformers import BertTokenizerFast, BertForSequenceClassification
+from transformers import BertTokenizerFast, BertForSequenceClassification,GPT2LMHeadModel,BartForConditionalGeneration
 import torch
+import math
+
+class CHSentenceSmoothScorer():
+    def __init__(self) -> None:
+        super().__init__()
+      
+        self.tokenizer = BertTokenizerFast.from_pretrained(
+            "fnlp/bart-base-chinese")
+        self.model = BartForConditionalGeneration.from_pretrained(
+            "fnlp/bart-base-chinese")
+
+    def __call__(self, sentences):
+        input_ids = self.tokenizer.batch_encode_plus(
+            sentences, return_tensors='pt',
+            padding=True,
+            max_length=50,
+            truncation='longest_first'
+        )['input_ids']
+        logits = self.model(input_ids).logits
+
+        softmax = torch.softmax(logits, dim=-1)
+
+        out = []
+        for i, sentence in enumerate(sentences):
+            sent_token_ids = input_ids[i].tolist()
+            sent_token_ids = list(
+                filter(lambda x: x not in [self.tokenizer.pad_token_id], sent_token_ids))
+            ppl = 0.0
+            for j, token_id in enumerate(sent_token_ids):
+                ppl += math.log(softmax[i][j][token_id].item())
+            ppl = -1*(ppl/len(sent_token_ids))
+            prob_socre = math.exp(ppl*-1)
+            out.append(prob_socre)
+
+        return out
+
 
 model = BertForSequenceClassification.from_pretrained('./ch-sent-check-model')
 tokenizer = BertTokenizerFast.from_pretrained('./ch-sent-check-model')
+smooth_scorer = CHSentenceSmoothScorer()
 
 def judge(sentence):
     input_ids = tokenizer(sentence,return_tensors='pt')['input_ids']
     out = model(input_ids)
     logits = out.logits
-    pred = torch.argmax(logits,dim=-1).item()
+    prob = torch.softmax(logits,dim=-1)
+    pred = torch.argmax(prob,dim=-1).item()
     pred_text = 'Incorrect' if pred == 0 else 'Correct'
-    return pred_text
+
+    correct_prob = prob[0][1].item()
+    pred_text = pred_text + f", score: {round(correct_prob*100,2)}"
+    smooth_score = round(smooth_scorer([sentence])[0]*100,2)
+    return pred_text,smooth_score
 
 iface = gr.Interface(
     fn=judge, 
     inputs=gr.Textbox(
         label="請輸入一段中文句子來檢測正確性",
         lines=1,
-        value="請注意用字的鄭確性",
     ), 
-    outputs="text"
+    outputs=[
+        gr.Textbox(
+            label="正確性檢查",
+            lines=1
+        ),
+        gr.Textbox(
+            label="流暢性檢查",
+            lines=1
+        )
+    ],
+    examples = [
+        '請注意用字的鄭確性',
+        '請注意用字的正確性'
+    ]
 )
 iface.launch()