DuongTrongChi/Toxic_Comment_Classification

1import gradio as gr2from transformers import AutoModelForSequenceClassification, AutoTokenizer, pipeline3 4 5class ToxicCommentClassification:6    def __init__(self, model_name: str):7        self.model = AutoModelForSequenceClassification.from_pretrained(model_name)8        self.tokenizer = AutoTokenizer.from_pretrained(model_name)9        self.pipeline = pipeline(10            "text-classification",11            model=self.model,12            tokenizer=self.tokenizer,13            return_all_scores=True,14        )15 16    def predict(self, text):17        res = self.pipeline(text)[0]18        results = dict()19        is_normal = True20        for x in res:21            results[x['label']] = x['score']22 23            if float(x['score']) > 0.8:24                is_normal = False25 26        # if is_normal:27        #     results['normal'] = 128        return results29 30 31def main():32    model = ToxicCommentClassification("DuongTrongChi/d-filter-v1.1")33    iface = gr.Interface(34        fn=model.predict,35        inputs=gr.Textbox(36            lines=3,37            placeholder="Hãy nhập nội dung vào đây",38            label="Input Text",39        ),40        outputs="label",41        title="Toxic Comment Classification",42        examples=[43            "Ôi chú chó này nhìn dễ thương thế!",44            "Cái lúc óc chó sống làm chi cho chật đất",45            "Cầm con dao này và đâm chết con chó này đi!",46            "Tôi dắt con mèo cưng của tôi đi đạo phố ở công viên."47        ],48    )49 50    iface.launch()51 52 53if __name__ == "__main__":54    main()