Update app.py

app.py

@@ -6,56 +6,90 @@ import gradio as gr
 
 from nltk import sent_tokenize
 from transformers import pipeline
+from gradio.themes.utils import red, green
+
+detector = pipeline(task='text-classification', model='SJTU-CL/RoBERTa-large-ArguGPT-sent')
+
+color_map = {
+    '0%': green.c400,
+    '10%': green.c300,
+    '20%': green.c200,
+    '30%': green.c100,
+    '40%': green.c50,
+    '50%': None,
+    '60%': red.c50,
+    '70%': red.c100,
+    '80%': red.c200,
+    '90%': red.c300,
+    '100%': red.c400
+}
 
-detector = pipeline(task='text-classification', model='yikang0131/argugpt-detector-sent')
 
 def predict_doc(doc):
     sents = sent_tokenize(doc)
     data = {'sentence': [], 'label': [], 'score': []}
     res = []
     for sent in sents:
-        label, score = predict_one_sent(sent)
+        prob = predict_one_sent(sent)
         data['sentence'].append(sent)
-        data['score'].append(score)
-        if label == 'LABEL_0':
-            res.append((sent, 'Human'))
+        data['score'].append(round(prob, 4))
+        if prob < 0.1: label = '0%'
+        elif prob < 0.2: label = '10%'
+        elif prob < 0.3: label = '20%',
+        elif prob < 0.4: label = '30%',
+        elif prob < 0.5: label = '40%',
+        elif prob < 0.6: label = '50%',
+        elif prob < 0.7: label = '60%',
+        elif prob < 0.8: label = '70%',
+        elif prob < 0.9: label = '80%',
+        elif prob < 1: label = '90%'
+        else: label = '100%'
+        res.append((sent, label))
+        if prob <= 0.5:
             data['label'].append('Human')
-        else: 
-            res.append((sent, 'Machine'))
-            data['label'].append('Machine')
+        else: data['label'].append('Machine')
     df = pd.DataFrame(data)
     df.to_csv('result.csv')
-    return res, df, 'result.csv'
+    overall_score = df.score.mean()
+    sum_str = ''
+    if overall_score <= 0.5: overall_label = 'Human'
+    else: overall_label = 'Machine'
+    sum_str = f'The essay is probably written by {overall_label}. The probability of being generated by AI is {overall_score}'
+    # print(sum_str)
+    return sum_str, res, df, 'result.csv'
 
 
 def predict_one_sent(sent):
+    '''
+    convert to prob 
+    LABEL_1, 0.66 -> 0.66
+    LABEL_0, 0.66 -> 0.34
+    '''
     res = detector(sent)[0]
-    return res['label'], res['score']
-
-
-iface = gr.Interface(
-    fn=predict_doc,
-    inputs=[
-        gr.Textbox(
-            label='Essay input',
-            info="Please enter essay in the textbox",
-            lines=5
-        )
-    ],
-    outputs=[
-        gr.HighlightedText(
-            label='Labeled Result',
-            show_legend=True
-        ).style(color_map={'Machine': 'red', 'Human': 'green'}),
-        gr.DataFrame(
-            label='Table with Probability Score',
-            max_rows=10
-        ),
-        gr.File(
-            label='CSV file storing data with all sentences'
-        )
-    ],
-    theme=gr.themes.Base()
-)
-
-iface.launch()
+    org_label, prob = res['label'], res['score']
+    if org_label == 'LABEL_0': prob = 1 - prob
+    return prob
+
+
+with gr.Blocks() as demo:
+    text_in = gr.Textbox(
+        lines=5, 
+        label='Essay input', 
+        info='Please enter the essay in the textbox'
+    )
+    button = gr.Button('Predict who writes this essay!')
+    summary = gr.Textbox(lines=1, label='Result summary')
+    sent_res = gr.HighlightedText(
+        label = 'Labeled Result'
+    ).style(color_map=color_map)
+    tab = gr.DataFrame(
+        label='Table with Probability Score',
+        max_rows=10
+    )
+    csv_f = gr.File(
+        label='CSV file storing data with all sentences.'
+    )
+    button.click(predict_doc, inputs=[text_in], outputs=[summary, sent_res, tab, csv_f], api_name='predict_doc')
+
+demo.launch()
+