app.py

import gradio as gr
from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline


class ToxicCommentClassification:
    def __init__(self, model_name: str):
        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.pipeline = pipeline(
            "text-classification",
            model=self.model,
            tokenizer=self.tokenizer,
            return_all_scores=True,
        )

    def predict(self, text):
        res = self.pipeline(text)[0]
        results = dict()
        is_normal = True
        for x in res:
            results[x['label']] = x['score']

            if float(x['score']) > 0.8:
                is_normal = False

        if is_normal:
            results['normal'] = 1
        return results


def main():
    model = ToxicCommentClassification("DuongTrongChi/facebook-commet-classification-filter-v4")
    iface = gr.Interface(
        fn=model.predict,
        inputs=gr.Textbox(
            lines=3,
            placeholder="Hãy nhập nội dung vào đây",
            label="Input Text",
        ),
        outputs="label",
        title="Toxic Comment Classification",
        examples=[
            "Ôi chú chó này nhìn dễ thương thế!",
            "Cái lúc óc chó sống làm chi cho chật đất",
            "Cầm con dao này và đâm chết con chó này đi!",
            "Tôi dắt con mèo cưng của tôi đi đạo phố ở công viên.",
            "Việt Nam là một nơi bạo loạn, người dân thì cướp bốc giết người. Quan chức nhà nước thì toàn bốc lộc nhân dân."
        ],
    )

    iface.launch()


if __name__ == "__main__":
    main()