Leetmonkey In Action. Darn LeetMonkey these days

app.py

@@ -7,6 +7,7 @@ import random
 import logging
 import os
 import autopep8
+import textwrap
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -14,7 +15,7 @@ logger = logging.getLogger(__name__)
 
 # Define the model options
 gguf_models = {
-    "Q8_0": "leetmonkey_peft__q8_0.gguf",
+    "Q8_0 (8-bit)": "leetmonkey_peft__q8_0.gguf",
     "Exact Copy": "leetmonkey_peft_exact_copy.gguf",
     "F16": "leetmonkey_peft_f16.gguf",
     "Super Block Q6": "leetmonkey_peft_super_block_q6.gguf"
@@ -32,10 +33,8 @@ def download_model(model_name):
     logger.info(f"Model downloaded: {model_path}")
     return model_path
 
-# Download the 8-bit model at startup
-q8_model_path = download_model(gguf_models["Q8_0"])
-
-# Load the 8-bit model
+# Download and load the 8-bit model at startup
+q8_model_path = download_model(gguf_models["Q8_0 (8-bit)"])
 llm = Llama(
     model_path=q8_model_path,
     n_ctx=2048,
@@ -43,12 +42,11 @@ llm = Llama(
     n_gpu_layers=0,
     verbose=False
 )
-
 logger.info("8-bit model loaded successfully")
 
 # Load the dataset
 dataset = load_dataset("sugiv/leetmonkey_python_dataset")
-val_dataset = dataset["train"].train_test_split(test_size=0.1)["test"]
+train_dataset = dataset["train"]
 
 # Generation parameters
 generation_kwargs = {
@@ -114,13 +112,18 @@ def extract_and_format_code(text):
     except:
         return formatted_code
 
-def update_problem():
-    sample = random.choice(val_dataset)
-    return sample['instruction']
+def select_random_problem():
+    return random.choice(train_dataset)['instruction']
 
-def update_solution(problem):
-    logger.info("Generating solution")
-    generated_output = generate_solution(problem)
+def update_solution(problem, model_name):
+    if model_name == "Q8_0 (8-bit)":
+        model = llm
+    else:
+        model_path = download_model(gguf_models[model_name])
+        model = Llama(model_path=model_path, n_ctx=2048, n_threads=4, n_gpu_layers=0, verbose=False)
+    
+    logger.info(f"Generating solution using {model_name} model")
+    generated_output = generate_solution(problem, model)
     formatted_code = extract_and_format_code(generated_output)
     logger.info("Solution generated successfully")
     return formatted_code
@@ -134,12 +137,13 @@ with gr.Blocks() as demo:
             select_problem_btn = gr.Button("Select Random Problem")
         
         with gr.Column():
+            model_dropdown = gr.Dropdown(choices=list(gguf_models.keys()), label="Select GGUF Model", value="Q8_0 (8-bit)")
             solution_display = gr.Code(label="Generated Solution", language="python", lines=25)
             generate_btn = gr.Button("Generate Solution")
     
-    select_problem_btn.click(update_problem, outputs=problem_display)
-    generate_btn.click(update_solution, inputs=[problem_display], outputs=solution_display)
+    select_problem_btn.click(select_random_problem, outputs=problem_display)
+    generate_btn.click(update_solution, inputs=[problem_display, model_dropdown], outputs=solution_display)
 
 if __name__ == "__main__":
     logger.info("Starting Gradio interface")
-    demo.launch(share=True)
+    demo.launch(share=True)