First version with APIs

app.py

@@ -1,28 +1,40 @@
-import gradio as gr
-from huggingface_hub import hf_hub_download
-from llama_cpp import Llama
+import os
 import re
-from datasets import load_dataset
-import random
 import logging
-import os
-import autopep8
 import textwrap
+import autopep8
+import gradio as gr
+from huggingface_hub import hf_hub_download
+from llama_cpp import Llama
 import jwt
-from datetime import datetime, timedelta
+from typing import Dict, Any
+import datetime
 
 # Set up logging
 logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # JWT settings
-JWT_SECRET = os.environ.get("JWT_SECRET", "your-secret-key")
+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": 512,
+    "stop": ["```", "### Instruction:", "### Response:"],
+    "echo": False,
+    "temperature": 0.05,
+    "top_k": 10,
+    "top_p": 0.9,
+    "repeat_penalty": 1.1
+}
+
 def download_model(model_name):
     logger.info(f"Downloading model: {model_name}")
     model_path = hf_hub_download(
@@ -48,22 +60,7 @@ llm = Llama(
 )
 logger.info("8-bit model loaded successfully")
 
-# Load the dataset
-dataset = load_dataset("sugiv/leetmonkey_python_dataset")
-train_dataset = dataset["train"]
-
-# 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):
+def generate_solution(instruction: str) -> str:
     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}
@@ -78,32 +75,27 @@ Here's the complete Python function implementation:
 ```python
 """
     
-    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
-        yield chunk["choices"][0]["text"]
+    response = llm(full_prompt, **generation_kwargs)
+    return response["choices"][0]["text"]
 
-def extract_and_format_code(text):
-    # Extract code between triple backticks
+def extract_and_format_code(text: str) -> str:
     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
+            indented_lines.append(line)
+        elif line.strip():
+            indented_lines.append('    ' + line)
         else:
-            indented_lines.append(line)  # Keep empty lines as is
+            indented_lines.append(line)
 
     formatted_code = '\n'.join(indented_lines)
 
@@ -112,66 +104,69 @@ def extract_and_format_code(text):
     except:
         return formatted_code
 
-def select_random_problem():
-    return random.choice(train_dataset)['instruction']
-
-def stream_solution(problem):
-    logger.info("Generating solution")
-    generated_text = ""
-    for token in generate_solution(problem):
-        generated_text += token
-        yield generated_text
-    
-    formatted_code = extract_and_format_code(generated_text)
-    logger.info("Solution generated successfully")
-    yield formatted_code
-
-def verify_token(token):
+def verify_token(token: str) -> bool:
     try:
         jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
         return True
-    except:
+    except jwt.PyJWTError:
         return False
 
-def generate_token():
-    expiration = datetime.utcnow() + timedelta(hours=1)
-    return jwt.encode({"exp": expiration}, JWT_SECRET, algorithm=JWT_ALGORITHM)
-
-def api_random_problem(token):
+def api_generate_solution(instruction: str, token: str) -> Dict[str, Any]:
     if not verify_token(token):
         return {"error": "Invalid token"}
-    return {"problem": select_random_problem()}
-
-def api_generate_solution(problem, token):
-    if not verify_token(token):
-        return {"error": "Invalid token"}
-    solution = "".join(list(stream_solution(problem)))
-    return {"solution": solution}
+    
+    generated_output = generate_solution(instruction)
+    formatted_code = extract_and_format_code(generated_output)
+    return {"solution": formatted_code}
 
-def api_explain_solution(solution, token):
+def api_explain_solution(code: 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{solution}\n\nExplanation:"
+    
+    explanation_prompt = f"Explain the following Python code:\n\n{code}\n\nExplanation:"
     explanation = llm(explanation_prompt, max_tokens=256)["choices"][0]["text"]
     return {"explanation": explanation}
 
-iface = gr.Interface(
-    fn=[api_random_problem, api_generate_solution, api_explain_solution, generate_token],
+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
+
+# Gradio interfaces
+iface_generate = gr.Interface(
+    fn=api_generate_solution,
     inputs=[
-        gr.Textbox(label="JWT Token"),
-        gr.Textbox(label="Problem"),
-        gr.Textbox(label="Solution")
+        gr.Textbox(label="LeetCode Problem Instruction"),
+        gr.Textbox(label="JWT Token")
     ],
-    outputs=[
-        gr.JSON(label="Random Problem"),
-        gr.JSON(label="Generated Solution"),
-        gr.JSON(label="Explanation"),
-        gr.Textbox(label="New 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")
     ],
-    title="LeetCode Problem Solver API",
-    description="API endpoints for generating and explaining LeetCode solutions."
+    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."
 )
 
+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")
-    iface.launch(share=True)
+    demo.launch(share=True)