app.py

import gradio as gr
from huggingface_hub import hf_hub_download
from llama_cpp import Llama
import re
import random
import logging
import os
import jwt
from typing import Dict, Any
import autopep8
import textwrap

from datasets import load_dataset
import random
import asyncio

# Set up logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

# 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)
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"

# Load the model
model_path = hf_hub_download(repo_id=REPO_ID, filename=MODEL_NAME, cache_dir="./models")
llm = Llama(model_path=model_path, n_ctx=1024, n_threads=8, n_gpu_layers=-1, verbose=False, mlock=True)
logger.info("8-bit model loaded successfully")

user_data = {}

# 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 verify_token(token: str) -> bool:
    try:
        jwt.decode(token, JWT_SECRET, algorithms=[JWT_ALGORITHM])
        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_explanation(problem: str, solution: str, token: str) -> Dict[str, Any]:
    if not verify_token(token):
        return {"error": "Invalid token"}
    
    system_prompt = "You are a Python coding assistant specialized in explaining LeetCode problem solutions. Provide a clear and concise explanation of the given solution."
    full_prompt = f"""### Instruction:
{system_prompt}

Problem:
{problem}

Solution:
{solution}

Explain this solution step by step.

### Response:
Here's the explanation of the solution:

"""
    
    generated_text = ""
    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
        generated_text += chunk["choices"][0]["text"]
    
    return {"explanation": generated_text}


def generate_solution(instruction: str, token: str) -> Dict[str, Any]:
    if not verify_token(token):
        return {"error": "Invalid token"}
    
    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:

{instruction}

### Response:
Here's the complete Python function implementation:

```python
"""
    
    generated_text = ""
    for chunk in llm(full_prompt, stream=True, **generation_kwargs):
        generated_text += chunk["choices"][0]["text"]
    
    formatted_code = extract_and_format_code(generated_text)
    user_data[token] = {"problem": instruction, "solution": formatted_code}
    return {"solution": formatted_code}

def random_problem(token: str) -> Dict[str, Any]:
    if not verify_token(token):
        return {"error": "Invalid token"}
    
    # 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']
    user_data[token] = {"problem": problem, "solution": None}
    
    return {"problem": problem}

def explain_solution(token: str) -> Dict[str, Any]:
    if not verify_token(token):
        return {"error": "Invalid token"}
    
    if token not in user_data or not user_data[token].get("solution"):
        return {"error": "No solution available to explain. Please generate a solution first."}
    
    problem = user_data[token]["problem"]
    solution = user_data[token]["solution"]
    
    return generate_explanation(problem, solution, token)

# Create Gradio interfaces
generate_interface = gr.Interface(
    fn=generate_solution,
    inputs=[gr.Textbox(label="Problem Instruction"), gr.Textbox(label="JWT Token")],
    outputs=gr.JSON(),
    title="Generate Solution API",
    description="Provide a LeetCode problem instruction and a valid JWT token to generate a solution."
)

random_problem_interface = gr.Interface(
    fn=random_problem,
    inputs=gr.Textbox(label="JWT Token"),
    outputs=gr.JSON(),
    title="Random Problem API",
    description="Provide a valid JWT token to get a random LeetCode problem."
)

explain_interface = gr.Interface(
    fn=explain_solution,
    inputs=gr.Textbox(label="JWT Token"),
    outputs=gr.JSON(),
    title="Explain Solution API",
    description="Provide a valid JWT token to get an explanation of the last generated solution."
)

demo = gr.TabbedInterface(
    [generate_interface, explain_interface, random_problem_interface],
    ["Generate Solution", "Explain Solution", "Random Problem"]
)

# Launch the Gradio app
demo.launch()