Leetmonkey In Action via Inference

app.py

@@ -13,18 +13,14 @@ import textwrap
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# Define the model options
-gguf_models = {
-    "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"
-}
+# Model settings
+MODEL_NAME = "leetmonkey_peft__q8_0.gguf"
+REPO_ID = "sugiv/leetmonkey-peft-gguf"
 
 def download_model(model_name):
     logger.info(f"Downloading model: {model_name}")
     model_path = hf_hub_download(
-        repo_id="sugiv/leetmonkey-peft-gguf",
+        repo_id=REPO_ID,
         filename=model_name,
         cache_dir="./models",
         force_download=True,
@@ -34,12 +30,12 @@ def download_model(model_name):
     return model_path
 
 # Download and load the 8-bit model at startup
-q8_model_path = download_model(gguf_models["Q8_0 (8-bit)"])
+model_path = download_model(MODEL_NAME)
 llm = Llama(
-    model_path=q8_model_path,
+    model_path=model_path,
     n_ctx=2048,
     n_threads=4,
-    n_gpu_layers=0,
+    n_gpu_layers=-1,  # Use all available GPU layers
     verbose=False
 )
 logger.info("8-bit model loaded successfully")
@@ -59,7 +55,7 @@ generation_kwargs = {
     "repeat_penalty": 1.1
 }
 
-def generate_solution(instruction, model):
+def generate_solution(instruction):
     system_prompt = "You are a Python coding assistant specialized in solving LeetCode problems. Provide only the complete implementation of the given function. Ensure proper indentation and formatting. Do not include any explanations or multiple solutions."
     full_prompt = f"""### Instruction:
 {system_prompt}
@@ -74,8 +70,8 @@ Here's the complete Python function implementation:
 ```python
 """
     
-    response = model(full_prompt, **generation_kwargs)
-    return response["choices"][0]["text"]
+    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
+        yield chunk["choices"][0]["text"]
 
 def extract_and_format_code(text):
     # Extract code between triple backticks
@@ -115,44 +111,10 @@ def extract_and_format_code(text):
 def select_random_problem():
     return random.choice(train_dataset)['instruction']
 
-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
-
-def stream_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")
-    system_prompt = "You are a Python coding assistant specialized in solving LeetCode problems. Provide only the complete implementation of the given function. Ensure proper indentation and formatting. Do not include any explanations or multiple solutions."
-    full_prompt = f"""### Instruction:
-{system_prompt}
-
-Implement the following function for the LeetCode problem:
-
-{problem}
-
-### Response:
-Here's the complete Python function implementation:
-
-```python
-"""
-    
+def stream_solution(problem):
+    logger.info("Generating solution")
     generated_text = ""
-    for chunk in model(full_prompt, stream=True, **generation_kwargs):
-        token = chunk["choices"][0]["text"]
+    for token in generate_solution(problem):
         generated_text += token
         yield generated_text
     
@@ -161,7 +123,7 @@ Here's the complete Python function implementation:
     yield formatted_code
 
 with gr.Blocks() as demo:
-    gr.Markdown("# LeetCode Problem Solver")
+    gr.Markdown("# LeetCode Problem Solver (8-bit GGUF Model)")
     
     with gr.Row():
         with gr.Column():
@@ -169,12 +131,11 @@ 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(select_random_problem, outputs=problem_display)
-    generate_btn.click(stream_solution, inputs=[problem_display, model_dropdown], outputs=solution_display)
+    generate_btn.click(stream_solution, inputs=[problem_display], outputs=solution_display)
 
 if __name__ == "__main__":
     logger.info("Starting Gradio interface")