import gradio as gr
from gradio import FlaggingCallback
from gradio.components import IOComponent

from typing import List, Optional, Any

import os
import argilla as rg

class ArgillaLogger(FlaggingCallback):
    def __init__(self, api_url, api_key):
        rg.init(api_url=api_url, api_key=api_key)
        
    def setup(self, components: List[IOComponent], flagging_dir: str):
        pass
    def flag(
        self,
        flag_data: List[Any],
        flag_option: Optional[str] = None,
        flag_index: Optional[int] = None,
        username: Optional[str] = None,
    ) -> int:
        text = flag_data[0]
        inference = flag_data[1]
        prediction = [(pred["label"], pred["confidence"]) for pred in inference["confidences"]]
        rg.log(
            name="cyber_bullying_sentiment_feedback",
            records=rg.TextClassificationRecord(text=text, prediction=prediction)
        )


description = "Get sentiment if a comment is 'no cyber bullying' or 'cyber bullying'.The predicted outputs refer to 'cyber bullying' for racism and sexism comments and 'no cyber bullying' for others comments.please do flag if you feel that the required output isn't desired output "
title = "Cyber Bullying Sentiment Classifier"
examples = [["lets have a coffee later on"],["go to hell Muslims "],["I respect your work ethics"]]
            
gr.Interface.load("models/Amitesh007/tw-sentiment-finetuned",
                 description = description,
                 title = title,
                 examples = examples,
                 allow_flagging="manual",
                
                 flagging_callback=ArgillaLogger(api_url="https://amitesh007-amitesh-argilla.hf.space", api_key="team.apikey")
                 ).launch()