| 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() |