from transformers import TextClassificationPipeline, AutoTokenizer, AutoModelForSequenceClassification
import gradio as gr
import os


class MaskedInterface:
    def __init__(self):
        self.models = [
        "Overfit-GM/bert-base-turkish-cased-offensive-mlm",
        "Overfit-GM/bert-base-turkish-uncased-offensive-mlm",
        "Overfit-GM/bert-base-turkish-128k-cased-offensive-mlm",
        "Overfit-GM/bert-base-turkish-128k-uncased-offensive-mlm",
        "Overfit-GM/convbert-base-turkish-mc4-cased-offensive-mlm",
        "Overfit-GM/convbert-base-turkish-mc4-uncased-offensive-mlm",
        "Overfit-GM/convbert-base-turkish-cased-offensive-mlm",
        "Overfit-GM/distilbert-base-turkish-cased-offensive-mlm",
        "Overfit-GM/electra-base-turkish-cased-discriminator-offensive-mlm",
        "Overfit-GM/electra-base-turkish-mc4-cased-discriminator-offensive-mlm",
        "Overfit-GM/electra-base-turkish-mc4-uncased-discriminator-offensive-mlm",
        "Overfit-GM/xlm-roberta-large-offensive-mlm",
        "Overfit-GM/mdeberta-v3-base-offensive-mlm"
    ]
        self.model_box = [
        gr.load(self.models[0], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[1], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[2], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[3], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[4], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[5], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[6], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[7], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[8], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[9], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[10], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[11], src='models', hf_token=os.environ['API_KEY']),
        gr.load(self.models[12], src='models', hf_token=os.environ['API_KEY'])
        ]

    def sentiment_analysis(self, text, model_choice):
        model = self.model_box[model_choice]
        output = model(text)
        return output
    
    def __call__(self):
        with gr.Blocks() as masked_interface:
            gr.HTML("""<h1 style="font-weight:600;font-size:50;margin-top:4px;margin-bottom:4px;text-align:center;">No Offense Fill Masks</h1></div>""")
            with gr.Row():
                with gr.Column():
                    model_choice = gr.Dropdown(label="Select Model", choices=[m for m in self.models], type="index", interactive=True)
                    input_text = gr.Textbox(label="Input", placeholder="senin ben [MASK]")
                    the_button = gr.Button(label="Run")
                with gr.Column():
                    output_window = gr.Label(num_top_classes=5)

            the_button.click(self.sentiment_analysis, inputs=[input_text, model_choice], outputs=[output_window])
            examples = gr.Examples(examples=["sen tam bir [MASK]", "erkekler [MASK] üstündür"],
                                inputs=[input_text])
            
        return masked_interface