Update app.py

app.py

@@ -1,48 +1,60 @@
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import load_dataset
 
-# Load the models
+# 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": gr.Interface.load("bigcode/python-stack-v1-functions-filtered-sc2-subset"),
-    "bigcode/python-stack-v1-functions-filtered-sc2": gr.Interface.load("bigcode/python-stack-v1-functions-filtered-sc2"),
-    "muellerzr/python-stack-v1-functions-filtered-llama-3-8B": gr.Interface.load("muellerzr/python-stack-v1-functions-filtered-llama-3-8B"),
-    "YoLo2000/python-stack-functions-filtered": gr.Interface.load("YoLo2000/python-stack-functions-filtered"),
-    "YoLo2000/python-stack-functions-filteredbigcode/python-stack-v1-functions-filtered-sc2": gr.Interface.load("YoLo2000/python-stack-functions-filteredbigcode/python-stack-v1-functions-filtered-sc2"),
-    "TheBloke/Python-Code-13B-GGUF": gr.Interface.load("TheBloke/Python-Code-13B-GGUF"),
-    "replit/replit-code-v1_5-3b": gr.Interface.load("replit/replit-code-v1_5-3b"),
-    "neulab/codebert-python": gr.Interface.load("neulab/codebert-python")
+    "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 the datasets
+# Load selected datasets
 datasets = {
-    "kye/all-huggingface-python-code": gr.Dataset.load("kye/all-huggingface-python-code"),
-    "ajibawa-2023/WikiHow": gr.Dataset.load("ajibawa-2023/WikiHow"),
-    "ajibawa-2023/Code-74k-ShareGPT": gr.Dataset.load("ajibawa-2023/Code-74k-ShareGPT"),
-    "ajibawa-2023/Software-Architectural-Frameworks": gr.Dataset.load("ajibawa-2023/Software-Architectural-Frameworks"),
-    "ajibawa-2023/Python-Code-23k-ShareGPT": gr.Dataset.load("ajibawa-2023/Python-Code-23k-ShareGPT"),
-    "HuggingFaceFW/fineweb": gr.Dataset.load("HuggingFaceFW/fineweb"),
-    "kye/all-huggingface-python-code-2": gr.Dataset.load("kye/all-huggingface-python-code-2"),
-    "suvadityamuk/huggingface-transformers-code-dataset": gr.Dataset.load("suvadityamuk/huggingface-transformers-code-dataset")
+    "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 interface
-def generate_code(prompt, model, dataset, temperature, max_length):
-    model_instance = models[model]
-    dataset_instance = datasets[dataset]
-    output = model_instance.generate(prompt, dataset_instance, temperature, max_length)
-    return output
+# 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,_default=0.5),
-________gr.Slider(label="Max_Length",_minimum=10,_maximum=1000,_default=200)
-____],
-    outputs="text"
+        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="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 interface
+# Launch the Gradio App
 iface.launch()
-