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/d-filter-v1.4")
    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."
        ],
    )

    iface.launch()


if __name__ == "__main__":
    main()