app.py

import nltk
nltk.download('punkt')

import pandas as pd
import gradio as gr

from nltk import sent_tokenize
from transformers import pipeline

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)
        data['sentence'].append(sent)
        data['score'].append(score)
        if label == 'LABEL_0':
            res.append((sent, 'Human'))
            data['label'].append('Human')
        else: 
            res.append((sent, 'Machine'))
            data['label'].append('Machine')
    df = pd.DataFrame(data)
    df.to_csv('result.csv')
    return res, df, 'result.csv'


def predict_one_sent(sent):
    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()