Fixing generate and random problem API

app.py

@@ -7,6 +7,14 @@ import logging
 import os
 import jwt
 from typing import Dict, Any
+import autopep8
+
+from datasets import load_dataset
+import random
+
+# Load the dataset (you might want to do this once at the start of your script)
+dataset = load_dataset("sugiv/leetmonkey_python_dataset")
+train_dataset = dataset["train"]
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
@@ -44,6 +52,38 @@ def verify_token(token: str) -> bool:
         return True
     except jwt.PyJWTError:
         return False
+    
+
+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)  # 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)  # Keep empty lines as is
+
+    formatted_code = '\n'.join(indented_lines)
+
+    try:
+        return autopep8.fix_code(formatted_code)
+    except:
+        return formatted_code
 
 def generate_solution(instruction: str, token: str) -> Dict[str, Any]:
     if not verify_token(token):
@@ -63,17 +103,12 @@ Here's the complete Python function implementation:
 ```python
 """
     
-    response = llm(full_prompt, **generation_kwargs)
-    generated_text = response["choices"][0]["text"]
-    
-    # Extract and format code
-    code_match = re.search(r'```python\s*(.*?)\s*```', generated_text, re.DOTALL)
-    if code_match:
-        code = code_match.group(1)
-    else:
-        code = generated_text
+    generated_text = ""
+    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
+        generated_text += chunk["choices"]["text"]
     
-    return {"solution": code}
+    formatted_code = extract_and_format_code(generated_text)
+    return {"solution": formatted_code}
 
 def stream_solution(instruction: str, token: str) -> Dict[str, Any]:
     if not verify_token(token):
@@ -103,13 +138,13 @@ def random_problem(token: str) -> Dict[str, Any]:
     if not verify_token(token):
         return {"error": "Invalid token"}
     
-    # This is a placeholder. You should replace it with actual logic to fetch a random problem from your dataset.
-    problems = [
-        "Implement a function to reverse a linked list",
-        "Write a function to find the maximum subarray sum",
-        "Implement a function to check if a binary tree is balanced"
-    ]
-    return {"problem": random.choice(problems)}
+    # Select a random problem from the dataset
+    random_item = random.choice(train_dataset)
+    
+    # Extract the instruction (problem statement) from the randomly selected item
+    problem = random_item['instruction']
+    
+    return {"problem": problem}
 
 # Create Gradio interfaces for each endpoint
 generate_interface = gr.Interface(

requirements.txt

@@ -4,4 +4,5 @@ datasets
 transformers
 autopep8
 huggingface_hub
-pyjwt
+pyjwt
+autopep8