File size: 4,046 Bytes
9c3d676 f18de00 9c3d676 f18de00 0f5897f 9c3d676 e1b0723 9c3d676 3ec3dc0 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 9c3d676 f18de00 3f58090 f18de00 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 |
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
# 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
}
def download_model(model_name: str) -> str:
logger.info(f"Downloading model: {model_name}")
model_path = hf_hub_download(
repo_id=REPO_ID,
filename=model_name,
cache_dir="./models",
force_download=True,
resume_download=True
)
logger.info(f"Model downloaded: {model_path}")
return model_path
# Download and load the 8-bit model at startup
model_path = download_model(MODEL_NAME)
llm = Llama(
model_path=model_path,
n_ctx=2048,
n_threads=4,
n_gpu_layers=-1, # Use all available GPU layers
verbose=False
)
logger.info("8-bit model loaded successfully")
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}
Implement the following function for the LeetCode problem:
{instruction}
### Response:
Here's the complete Python function implementation:
```python
"""
response = llm(full_prompt, **generation_kwargs)
return response["choices"][0]["text"]
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
code = re.sub(r'^.*?(?=def\s+\w+\s*\()', '', code, flags=re.DOTALL)
code = textwrap.dedent(code)
lines = code.split('\n')
func_def_index = next((i for i, line in enumerate(lines) if line.strip().startswith('def ')), 0)
indented_lines = [lines[func_def_index]]
for line in lines[func_def_index + 1:]:
if line.strip():
indented_lines.append(' ' + line)
else:
indented_lines.append(line)
formatted_code = '\n'.join(indented_lines)
try:
return autopep8.fix_code(formatted_code)
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"}
logger.info("Generating solution")
generated_output = generate_solution(instruction)
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"]
]
)
if __name__ == "__main__":
api.launch(server_name="0.0.0.0", server_port=7860)
|