Gosh this Leetmonkey

app.py

@@ -1,14 +1,15 @@
-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 re
+from datasets import load_dataset
+import random
+import logging
+import os
+import autopep8
+import textwrap
 import jwt
 from typing import Dict, Any
-import datetime
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -24,16 +25,9 @@ JWT_ALGORITHM = "HS256"
 MODEL_NAME = "leetmonkey_peft__q8_0.gguf"
 REPO_ID = "sugiv/leetmonkey-peft-gguf"
 
-# Generation parameters
-generation_kwargs = {
-    "max_tokens": 512,
-    "stop": ["```", "### Instruction:", "### Response:"],
-    "echo": False,
-    "temperature": 0.05,
-    "top_k": 10,
-    "top_p": 0.9,
-    "repeat_penalty": 1.1
-}
+# Load the dataset
+dataset = load_dataset("sugiv/leetmonkey_python_dataset")
+train_dataset = dataset["train"]
 
 def download_model(model_name):
     logger.info(f"Downloading model: {model_name}")
@@ -60,7 +54,18 @@ llm = Llama(
 )
 logger.info("8-bit model loaded successfully")
 
-def generate_solution(instruction: str) -> str:
+# Generation parameters
+generation_kwargs = {
+    "max_tokens": 512,
+    "stop": ["```", "### Instruction:", "### Response:"],
+    "echo": False,
+    "temperature": 0.05,
+    "top_k": 10,
+    "top_p": 0.9,
+    "repeat_penalty": 1.1
+}
+
+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}
@@ -75,27 +80,32 @@ Here's the complete Python function implementation:
 ```python
 """
     
-    response = llm(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: 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
 
+    # Dedent the code to remove any common leading whitespace
     code = textwrap.dedent(code)
+
+    # Split the code into lines
     lines = code.split('\n')
 
+    # Ensure proper indentation
     indented_lines = []
     for line in lines:
         if line.strip().startswith('class') or line.strip().startswith('def'):
-            indented_lines.append(line)
-        elif line.strip():
-            indented_lines.append('    ' + line)
+            indented_lines.append(line)  # Keep class and function definitions as is
+        elif 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)
 
@@ -104,6 +114,9 @@ def extract_and_format_code(text: str) -> str:
     except:
         return formatted_code
 
+def select_random_problem():
+    return random.choice(train_dataset)['instruction']
+
 def verify_token(token: str) -> bool:
     try:
         jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
@@ -111,62 +124,67 @@ def verify_token(token: str) -> bool:
     except jwt.PyJWTError:
         return False
 
+last_generated_solution = ""
+
 def api_generate_solution(instruction: str, token: str) -> Dict[str, Any]:
     if not verify_token(token):
         return {"error": "Invalid token"}
     
-    generated_output = generate_solution(instruction)
-    formatted_code = extract_and_format_code(generated_output)
-    return {"solution": formatted_code}
+    global last_generated_solution
+    generated_text = "".join(list(generate_solution(instruction)))
+    last_generated_solution = extract_and_format_code(generated_text)
+    return {"solution": last_generated_solution}
 
-def api_explain_solution(code: str, token: str) -> Dict[str, Any]:
+def api_stream_solution(instruction: str, token: str) -> Dict[str, Any]:
     if not verify_token(token):
         return {"error": "Invalid token"}
     
-    explanation_prompt = f"Explain the following Python code:\n\n{code}\n\nExplanation:"
+    def generate():
+        global last_generated_solution
+        generated_text = ""
+        for token in generate_solution(instruction):
+            generated_text += token
+            yield {"token": token}
+        
+        last_generated_solution = extract_and_format_code(generated_text)
+        yield {"solution": last_generated_solution}
+    
+    return generate()
+
+def api_explain_solution(token: str) -> Dict[str, Any]:
+    if not verify_token(token):
+        return {"error": "Invalid token"}
+    
+    if not last_generated_solution:
+        return {"error": "No solution has been generated yet"}
+    
+    explanation_prompt = f"Explain the following Python code:\n\n{last_generated_solution}\n\nExplanation:"
     explanation = llm(explanation_prompt, max_tokens=256)["choices"][0]["text"]
     return {"explanation": explanation}
 
-def generate_token() -> str:
-    expiration = datetime.datetime.utcnow() + datetime.timedelta(hours=1)
-    payload = {"exp": expiration}
-    token = jwt.encode(payload, JWT_SECRET, algorithm=JWT_ALGORITHM)
-    return token
+def api_random_problem(token: str) -> Dict[str, Any]:
+    if not verify_token(token):
+        return {"error": "Invalid token"}
+    
+    return {"problem": select_random_problem()}
 
-# Gradio interfaces
-iface_generate = gr.Interface(
-    fn=api_generate_solution,
+# Gradio interface
+iface = gr.Interface(
+    fn=[api_generate_solution, api_stream_solution, api_explain_solution, api_random_problem],
     inputs=[
-        gr.Textbox(label="LeetCode Problem Instruction"),
+        gr.Textbox(label="LeetCode Problem"),
         gr.Textbox(label="JWT Token")
     ],
-    outputs=gr.JSON(label="Generated Solution"),
-    title="LeetCode Problem Solver API - Generate Solution",
-    description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution."
-)
-
-iface_explain = gr.Interface(
-    fn=api_explain_solution,
-    inputs=[
-        gr.Textbox(label="Code to Explain"),
-        gr.Textbox(label="JWT Token")
+    outputs=[
+        gr.JSON(label="Generated Solution"),
+        gr.JSON(label="Streamed Solution"),
+        gr.JSON(label="Explanation"),
+        gr.JSON(label="Random Problem")
     ],
-    outputs=gr.JSON(label="Explanation"),
-    title="LeetCode Problem Solver API - Explain Solution",
-    description="Provide a code snippet and a valid JWT token to get an explanation."
+    title="LeetCode Problem Solver API",
+    description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution, get an explanation, or retrieve a random problem."
 )
 
-iface_token = gr.Interface(
-    fn=generate_token,
-    inputs=[],
-    outputs=gr.Textbox(label="Generated JWT Token"),
-    title="Generate JWT Token",
-    description="Generate a new JWT token for API authentication."
-)
-
-# Combine interfaces
-demo = gr.TabbedInterface([iface_generate, iface_explain, iface_token], ["Generate Solution", "Explain Solution", "Generate Token"])
-
 if __name__ == "__main__":
     logger.info("Starting Gradio API")
-    demo.launch(share=True)
+    iface.launch(share=True)