Leetmonkey In Action via Inference

app.py

@@ -1,43 +1,30 @@
-import os
-import re
-import logging
-import textwrap
-import autopep8
 import gradio as gr
 from huggingface_hub import hf_hub_download
 from llama_cpp import Llama
-import jwt
-from typing import Dict, Any
+import re
+from datasets import load_dataset
+import random
+import logging
+import os
+import autopep8
+import textwrap
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
-# JWT settings
-JWT_SECRET = os.environ.get("JWT_SECRET")
-if not JWT_SECRET:
-    raise ValueError("JWT_SECRET environment variable is not set")
-JWT_ALGORITHM = "HS256"
-
-# Model settings
-MODEL_NAME = "leetmonkey_peft__q8_0.gguf"
-REPO_ID = "sugiv/leetmonkey-peft-gguf"
-
-# Generation parameters
-generation_kwargs = {
-    "max_tokens": 2048,
-    "stop": ["```", "### Instruction:", "### Response:"],
-    "echo": False,
-    "temperature": 0.2,
-    "top_k": 50,
-    "top_p": 0.95,
-    "repeat_penalty": 1.1
+# 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"
 }
 
-def download_model(model_name: str) -> str:
+def download_model(model_name):
     logger.info(f"Downloading model: {model_name}")
     model_path = hf_hub_download(
-        repo_id=REPO_ID,
+        repo_id="sugiv/leetmonkey-peft-gguf",
         filename=model_name,
         cache_dir="./models",
         force_download=True,
@@ -47,17 +34,32 @@ def download_model(model_name: str) -> str:
     return model_path
 
 # Download and load the 8-bit model at startup
-model_path = download_model(MODEL_NAME)
+q8_model_path = download_model(gguf_models["Q8_0 (8-bit)"])
 llm = Llama(
-    model_path=model_path,
+    model_path=q8_model_path,
     n_ctx=2048,
     n_threads=4,
-    n_gpu_layers=-1,  # Use all available GPU layers
+    n_gpu_layers=0,
     verbose=False
 )
 logger.info("8-bit model loaded successfully")
 
-def generate_solution(instruction: str) -> str:
+# Load the dataset
+dataset = load_dataset("sugiv/leetmonkey_python_dataset")
+train_dataset = dataset["train"]
+
+# Generation parameters
+generation_kwargs = {
+    "max_tokens": 2048,
+    "stop": ["```", "### Instruction:", "### Response:"],
+    "echo": False,
+    "temperature": 0.2,
+    "top_k": 50,
+    "top_p": 0.95,
+    "repeat_penalty": 1.1
+}
+
+def generate_solution(instruction, 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}
@@ -72,26 +74,36 @@ Here's the complete Python function implementation:
 ```python
 """
     
-    response = llm(full_prompt, **generation_kwargs)
+    response = model(full_prompt, **generation_kwargs)
     return response["choices"][0]["text"]
 
-def extract_and_format_code(text: str) -> str:
+def extract_and_format_code(text):
+    # Extract code between triple backticks
     code_match = re.search(r'```python\s*(.*?)\s*```', text, re.DOTALL)
     if code_match:
         code = code_match.group(1)
     else:
         code = text
 
+    # Remove any text before the function definition
     code = re.sub(r'^.*?(?=def\s+\w+\s*\()', '', code, flags=re.DOTALL)
+
+    # Dedent the code to remove any common leading whitespace
     code = textwrap.dedent(code)
+
+    # Split the code into lines
     lines = code.split('\n')
+
+    # Find the function definition line
     func_def_index = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), 0)
-    indented_lines = [lines[func_def_index]]
+
+    # Ensure proper indentation
+    indented_lines = [lines[func_def_index]]  # Keep the function definition as is
     for line in lines[func_def_index + 1:]:
-        if line.strip():
-            indented_lines.append('    ' + line)
+        if line.strip():  # If the line is not empty
+            indented_lines.append('    ' + line)  # Add 4 spaces of indentation
         else:
-            indented_lines.append(line)
+            indented_lines.append(line)  # Keep empty lines as is
 
     formatted_code = '\n'.join(indented_lines)
 
@@ -100,38 +112,70 @@ def extract_and_format_code(text: str) -> str:
     except:
         return formatted_code
 
-def verify_token(token: str) -> bool:
-    try:
-        jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
-        return True
-    except jwt.PyJWTError:
-        return False
-
-def generate_code(instruction: str, token: str) -> Dict[str, Any]:
-    if not verify_token(token):
-        return {"error": "Invalid token"}
+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("Generating solution")
-    generated_output = generate_solution(instruction)
+    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 {"solution": formatted_code}
-
-# Gradio API
-api = gr.Interface(
-    fn=generate_code,
-    inputs=[
-        gr.Textbox(label="LeetCode Problem Instruction"),
-        gr.Textbox(label="JWT Token")
-    ],
-    outputs=gr.JSON(),
-    title="LeetCode Problem Solver API",
-    description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution.",
-    examples=[
-        ["Implement a function to reverse a linked list", "your_jwt_token_here"],
-        ["Write a function to find the maximum subarray sum", "your_jwt_token_here"]
-    ]
-)
+    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
+"""
+    
+    generated_text = ""
+    for chunk in model(full_prompt, stream=True, **generation_kwargs):
+        token = chunk["choices"][0]["text"]
+        generated_text += token
+        yield generated_text
+    
+    formatted_code = extract_and_format_code(generated_text)
+    logger.info("Solution generated successfully")
+    yield formatted_code
+
+with gr.Blocks() as demo:
+    gr.Markdown("# LeetCode Problem Solver")
+    
+    with gr.Row():
+        with gr.Column():
+            problem_display = gr.Textbox(label="LeetCode Problem", lines=10)
+            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)
 
 if __name__ == "__main__":
-    api.launch(server_name="0.0.0.0", server_port=7860)
+    logger.info("Starting Gradio interface")
+    demo.launch(share=True)

requirements.txt

@@ -1,8 +1,6 @@
 gradio
 llama-cpp-python
-huggingface_hub
-pyjwt
+datasets
+transformers
 autopep8
-fastapi
-uvicorn
-pydantic
+huggingface_hub