import gradio as gr
from transformers import AutoModelForCausalLM, AutoTokenizer
from datasets import load_dataset

# Define model loading function
def load_model(model_name):
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForCausalLM.from_pretrained(model_name)
    return tokenizer, model

# Load selected models
models = {
    "bigcode/python-stack-v1-functions-filtered-sc2-subset": "bigcode/python-stack-v1-functions-filtered-sc2-subset",
    "bigcode/python-stack-v1-functions-filtered-sc2": "bigcode/python-stack-v1-functions-filtered-sc2",
    "muellerzr/python-stack-v1-functions-filtered-llama-3-8B": "muellerzr/python-stack-v1-functions-filtered-llama-3-8B",
    "TheBloke/Python-Code-13B-GGUF": "TheBloke/Python-Code-13B-GGUF",
    "replit/replit-code-v1_5-3b": "replit/replit-code-v1_5-3b",
    "neulab/codebert-python": "neulab/codebert-python"
}

# Load selected datasets
datasets = {
    "kye/all-huggingface-python-code": "kye/all-huggingface-python-code",
    "ajibawa-2023/Python-Code-23k-ShareGPT": "ajibawa-2023/Python-Code-23k-ShareGPT",
    "suvadityamuk/huggingface-transformers-code-dataset": "suvadityamuk/huggingface-transformers-code-dataset"
}

# Define the function for code generation
def generate_code(prompt, model_name, dataset_name, temperature, max_length):
    tokenizer, model = load_model(models[model_name])
    
    # Load dataset (for reference, not directly used)
    dataset = load_dataset(datasets[dataset_name], split="train")

    # Tokenize input prompt
    inputs = tokenizer(prompt, return_tensors="pt")
    
    # Generate output
    output_ids = model.generate(**inputs, temperature=temperature, max_length=max_length)
    generated_code = tokenizer.decode(output_ids[0], skip_special_tokens=True)
    
    return generated_code

# Create Gradio Interface
iface = gr.Interface(
    fn=generate_code,
    inputs=[
        gr.Textbox(label="Prompt"),
        gr.Dropdown(label="Model", choices=list(models.keys())),
        gr.Dropdown(label="Dataset", choices=list(datasets.keys())),
        gr.Slider(label="Temperature", minimum=0.1, maximum=1.0, value=0.5),
        gr.Slider(label="Max Length", minimum=10, maximum=1000, value=200)
    ],
    outputs="text",
    title="Ultimate Code Generator Python: AI Code Generator with Hugging Face Models",
    description="Select a model and dataset, input a prompt, and generate Python code using AI models."
)

# Launch the Gradio App
iface.launch()