import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification

models = ["bert-base-uncased", "roberta-base"]
datasets = ["None", 
            "vedantgaur/GPTOutputs-MWP - AI Data Only",
            "vedantgaur/GPTOutputs-MWP - Human Data Only",
            "vedantgaur/GPTOutputs-MWP - Both AI and Human Data",
            "dmitva/human_ai_generated_text - Both AI and Human Data"]

# Mapping of user-selected model and dataset to actual model name on Hugging Face
model_mapping = {
    ("bert-base-uncased", "None"): "bert-base-uncased",
    ("bert-base-uncased", "vedantgaur/GPTOutputs-MWP - AI Data Only"): "SkwarczynskiP/bert-base-uncased-finetuned-vedantgaur-AI-generated",
    ("bert-base-uncased", "vedantgaur/GPTOutputs-MWP - Human Data Only"): "SkwarczynskiP/bert-base-uncased-finetuned-vedantgaur-human-generated",
    ("bert-base-uncased", "vedantgaur/GPTOutputs-MWP - Both AI and Human Data"): "SkwarczynskiP/bert-base-uncased-finetuned-vedantgaur-AI-and-human-generated",
    ("bert-base-uncased", "dmitva/human_ai_generated_text"): "SkwarczynskiP/bert-base-uncased-finetuned-dmitva-AI-and-human-generated",
    ("roberta-base", "None"): "roberta-base",
    ("roberta-base", "vedantgaur/GPTOutputs-MWP - AI Data Only"): "SkwarczynskiP/roberta-base-finetuned-vedantgaur-AI-generated",
    ("roberta-base", "vedantgaur/GPTOutputs-MWP - Human Data Only"): "SkwarczynskiP/roberta-base-finetuned-vedantgaur-human-generated",
    ("roberta-base", "vedantgaur/GPTOutputs-MWP - Both AI and Human Data"): "SkwarczynskiP/roberta-base-finetuned-vedantgaur-AI-and-human-generated",
    ("roberta-base", "dmitva/human_ai_generated_text"): "SkwarczynskiP/roberta-base-finetuned-dmitva-AI-and-human-generated"
}

def detect_ai_generated_text(model: str, dataset: str, text: str) -> str:
    # Get the fine-tuned model using mapping
    finetuned_model = model_mapping.get((model, dataset))

    # Load the specific fine-tuned model
    tokenizer = AutoTokenizer.from_pretrained(finetuned_model)
    model = AutoModelForSequenceClassification.from_pretrained(finetuned_model)

    # Classify the input based on the fine-tuned model
    classifier = pipeline('text-classification', model=model, tokenizer=tokenizer)
    result = classifier(text)
    return "AI-generated" if result[0]['label'] == 'LABEL_1' else "Not AI-generated"

interface = gr.Interface(
    fn=detect_ai_generated_text,
    inputs=[
        gr.Dropdown(choices=models, label="Model"),
        gr.Dropdown(choices=datasets, label="Dataset"),
        gr.Textbox(lines=5, label="Input Text")
    ],
    outputs=gr.Textbox(label="Output")
)

if __name__ == "__main__":
    interface.launch()