Spaces:
Sleeping
Sleeping
Sushwetabm
commited on
Commit
Β·
f59cf24
0
Parent(s):
Deploy ML microservice to Hugging Face Space
Browse files- .dockerignore +12 -0
- .gitignore +9 -0
- Dockerfile +32 -0
- __init__.py +0 -0
- analyzer.py +326 -0
- app.py +55 -0
- main.py +409 -0
- model.py +124 -0
- requirements.txt +24 -0
- setup.py +106 -0
.dockerignore
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
__pycache__
|
2 |
+
.venv
|
3 |
+
.git
|
4 |
+
*.md
|
5 |
+
*.pdf
|
6 |
+
*.pt
|
7 |
+
*.bin
|
8 |
+
*.log
|
9 |
+
.venv/
|
10 |
+
*.pyc
|
11 |
+
.DS_Store
|
12 |
+
model_cache/
|
.gitignore
ADDED
@@ -0,0 +1,9 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
__pycache__/
|
2 |
+
*.pyc
|
3 |
+
venv/
|
4 |
+
.env
|
5 |
+
.venv
|
6 |
+
Lib
|
7 |
+
model_cache/
|
8 |
+
offload/
|
9 |
+
|
Dockerfile
ADDED
@@ -0,0 +1,32 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# β
Use official slim image
|
2 |
+
FROM python:3.10-slim
|
3 |
+
|
4 |
+
# β
Set working directory
|
5 |
+
WORKDIR /app
|
6 |
+
|
7 |
+
# β
Set environment variables early to ensure cache use in setup.py
|
8 |
+
ENV TRANSFORMERS_CACHE=/app/model_cache \
|
9 |
+
HF_HOME=/app/model_cache \
|
10 |
+
TORCH_HOME=/app/model_cache \
|
11 |
+
TOKENIZERS_PARALLELISM=false \
|
12 |
+
OMP_NUM_THREADS=4
|
13 |
+
|
14 |
+
# β
Install only necessary OS packages and clean cache
|
15 |
+
RUN apt-get update && apt-get install -y git \
|
16 |
+
&& rm -rf /var/lib/apt/lists/*
|
17 |
+
|
18 |
+
# β
Copy files (excluding model_cache, logs, etc. via .dockerignore)
|
19 |
+
COPY . .
|
20 |
+
|
21 |
+
# β
Upgrade pip + install deps without cache
|
22 |
+
RUN pip install --upgrade pip \
|
23 |
+
&& pip install --no-cache-dir -r requirements.txt
|
24 |
+
|
25 |
+
# β
Run setup.py to download the model
|
26 |
+
RUN python setup.py
|
27 |
+
|
28 |
+
# β
Expose Hugging Face Space-required port
|
29 |
+
EXPOSE 7860
|
30 |
+
|
31 |
+
# β
Launch FastAPI on port 7860 for HF Space
|
32 |
+
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "7860", "--timeout-]()
|
__init__.py
ADDED
File without changes
|
analyzer.py
ADDED
@@ -0,0 +1,326 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# import json
|
2 |
+
|
3 |
+
# def analyze_code(language, code, tokenizer, model):
|
4 |
+
# messages = [
|
5 |
+
# {
|
6 |
+
# "role": "system",
|
7 |
+
# "content": (
|
8 |
+
# "You are a helpful and expert-level AI code reviewer and bug fixer. "
|
9 |
+
# "Your task is to analyze the given buggy code in the specified programming language, "
|
10 |
+
# "identify bugs (logical, syntax, runtime, etc.), and fix them. "
|
11 |
+
# "Return a JSON object with the following keys:\n\n"
|
12 |
+
# "1. 'bug_analysis': a list of objects, each containing:\n"
|
13 |
+
# " - 'line_number': the line number (approximate if needed)\n"
|
14 |
+
# " - 'error_message': a short name of the bug\n"
|
15 |
+
# " - 'explanation': short explanation of the problem\n"
|
16 |
+
# " - 'fix_suggestion': how to fix it\n"
|
17 |
+
# "2. 'corrected_code': the entire corrected code block.\n\n"
|
18 |
+
# "Respond with ONLY the raw JSON object, no extra commentary or markdown."
|
19 |
+
# )
|
20 |
+
# },
|
21 |
+
# {
|
22 |
+
# "role": "user",
|
23 |
+
# "content": f"π» Language: {language}\nπ Buggy Code:\n```{language.lower()}\n{code.strip()}\n```"
|
24 |
+
# }
|
25 |
+
# ]
|
26 |
+
|
27 |
+
# inputs = tokenizer.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
|
28 |
+
# attention_mask = (inputs != tokenizer.pad_token_id).long()
|
29 |
+
|
30 |
+
# outputs = model.generate(
|
31 |
+
# inputs,
|
32 |
+
# attention_mask=attention_mask,
|
33 |
+
# max_new_tokens=1024,
|
34 |
+
# do_sample=False,
|
35 |
+
# pad_token_id=tokenizer.eos_token_id,
|
36 |
+
# eos_token_id=tokenizer.eos_token_id
|
37 |
+
# )
|
38 |
+
|
39 |
+
# response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
|
40 |
+
|
41 |
+
# # Try parsing response to JSON
|
42 |
+
# try:
|
43 |
+
# json_output = json.loads(response)
|
44 |
+
# return json_output
|
45 |
+
# except json.JSONDecodeError:
|
46 |
+
# print("β οΈ Could not decode response into JSON. Here's the raw output:\n")
|
47 |
+
# print(response)
|
48 |
+
# return None
|
49 |
+
# import json
|
50 |
+
# import logging
|
51 |
+
# import time
|
52 |
+
# import torch
|
53 |
+
|
54 |
+
# # Configure logging
|
55 |
+
# logger = logging.getLogger(__name__)
|
56 |
+
|
57 |
+
# def analyze_code(language, code, tokenizer, model):
|
58 |
+
# """
|
59 |
+
# Analyze code and return bug analysis with improved logging and error handling
|
60 |
+
# """
|
61 |
+
# start_time = time.time()
|
62 |
+
# logger.info(f"π Starting analysis for {language} code ({len(code)} characters)")
|
63 |
+
|
64 |
+
# try:
|
65 |
+
# # Prepare messages
|
66 |
+
# messages = [
|
67 |
+
# {
|
68 |
+
# "role": "system",
|
69 |
+
# "content": (
|
70 |
+
# "You are a helpful and expert-level AI code reviewer and bug fixer. "
|
71 |
+
# "Your task is to analyze the given buggy code in the specified programming language, "
|
72 |
+
# "identify bugs (logical, syntax, runtime, etc.), and fix them. "
|
73 |
+
# "Return a JSON object with the following keys:\n\n"
|
74 |
+
# "1. 'bug_analysis': a list of objects, each containing:\n"
|
75 |
+
# " - 'line_number': the line number (approximate if needed)\n"
|
76 |
+
# " - 'error_message': a short name of the bug\n"
|
77 |
+
# " - 'explanation': short explanation of the problem\n"
|
78 |
+
# " - 'fix_suggestion': how to fix it\n"
|
79 |
+
# "2. 'corrected_code': the entire corrected code block.\n\n"
|
80 |
+
# "Respond with ONLY the raw JSON object, no extra commentary or markdown."
|
81 |
+
# )
|
82 |
+
# },
|
83 |
+
# {
|
84 |
+
# "role": "user",
|
85 |
+
# "content": f"π» Language: {language}\nπ Buggy Code:\n```{language.lower()}\n{code.strip()}\n```"
|
86 |
+
# }
|
87 |
+
# ]
|
88 |
+
|
89 |
+
# logger.info("π§ Applying chat template...")
|
90 |
+
# inputs = tokenizer.apply_chat_template(
|
91 |
+
# messages,
|
92 |
+
# add_generation_prompt=True,
|
93 |
+
# return_tensors="pt"
|
94 |
+
# ).to(model.device)
|
95 |
+
|
96 |
+
# attention_mask = (inputs != tokenizer.pad_token_id).long()
|
97 |
+
|
98 |
+
# logger.info(f"π Input length: {inputs.shape[1]} tokens")
|
99 |
+
# logger.info("π Starting model generation...")
|
100 |
+
|
101 |
+
# generation_start = time.time()
|
102 |
+
|
103 |
+
# # Generate with more conservative settings
|
104 |
+
# with torch.no_grad(): # Ensure no gradients are computed
|
105 |
+
# outputs = model.generate(
|
106 |
+
# inputs,
|
107 |
+
# attention_mask=attention_mask,
|
108 |
+
# max_new_tokens=512, # Reduced from 1024 for faster inference
|
109 |
+
# do_sample=False,
|
110 |
+
# temperature=0.1, # Add temperature for more consistent output
|
111 |
+
# pad_token_id=tokenizer.eos_token_id,
|
112 |
+
# eos_token_id=tokenizer.eos_token_id,
|
113 |
+
# use_cache=True, # Enable KV cache for efficiency
|
114 |
+
# )
|
115 |
+
|
116 |
+
# generation_time = time.time() - generation_start
|
117 |
+
# logger.info(f"β‘ Generation completed in {generation_time:.2f} seconds")
|
118 |
+
|
119 |
+
# logger.info("π Decoding response...")
|
120 |
+
# response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
|
121 |
+
|
122 |
+
# logger.info(f"π Response length: {len(response)} characters")
|
123 |
+
# logger.info(f"π First 100 chars: {response[:100]}...")
|
124 |
+
|
125 |
+
# # Try parsing response to JSON
|
126 |
+
# logger.info("π Attempting to parse JSON...")
|
127 |
+
# try:
|
128 |
+
# # Clean up response - remove any markdown formatting
|
129 |
+
# cleaned_response = response.strip()
|
130 |
+
# if cleaned_response.startswith('```json'):
|
131 |
+
# cleaned_response = cleaned_response[7:]
|
132 |
+
# if cleaned_response.startswith('```'):
|
133 |
+
# cleaned_response = cleaned_response[3:]
|
134 |
+
# if cleaned_response.endswith('```'):
|
135 |
+
# cleaned_response = cleaned_response[:-3]
|
136 |
+
|
137 |
+
# cleaned_response = cleaned_response.strip()
|
138 |
+
|
139 |
+
# json_output = json.loads(cleaned_response)
|
140 |
+
|
141 |
+
# total_time = time.time() - start_time
|
142 |
+
# logger.info(f"β
Analysis completed successfully in {total_time:.2f} seconds")
|
143 |
+
|
144 |
+
# # Validate the JSON structure
|
145 |
+
# if not isinstance(json_output, dict):
|
146 |
+
# raise ValueError("Response is not a dictionary")
|
147 |
+
|
148 |
+
# if 'bug_analysis' not in json_output:
|
149 |
+
# logger.warning("β οΈ Missing 'bug_analysis' key, adding empty list")
|
150 |
+
# json_output['bug_analysis'] = []
|
151 |
+
|
152 |
+
# if 'corrected_code' not in json_output:
|
153 |
+
# logger.warning("β οΈ Missing 'corrected_code' key, adding original code")
|
154 |
+
# json_output['corrected_code'] = code
|
155 |
+
|
156 |
+
# return json_output
|
157 |
+
|
158 |
+
# except json.JSONDecodeError as e:
|
159 |
+
# logger.error(f"β JSON decode error: {e}")
|
160 |
+
# logger.error(f"π Raw response: {repr(response)}")
|
161 |
+
|
162 |
+
# # Return a fallback structure with the raw response
|
163 |
+
# fallback_response = {
|
164 |
+
# "bug_analysis": [{
|
165 |
+
# "line_number": 1,
|
166 |
+
# "error_message": "Analysis parsing failed",
|
167 |
+
# "explanation": "The AI model returned a response that couldn't be parsed as JSON",
|
168 |
+
# "fix_suggestion": "Please try again or check the code format"
|
169 |
+
# }],
|
170 |
+
# "corrected_code": code,
|
171 |
+
# "raw_output": response,
|
172 |
+
# "parsing_error": str(e)
|
173 |
+
# }
|
174 |
+
|
175 |
+
# return fallback_response
|
176 |
+
|
177 |
+
# except Exception as e:
|
178 |
+
# total_time = time.time() - start_time
|
179 |
+
# logger.error(f"β Analysis failed after {total_time:.2f} seconds: {str(e)}")
|
180 |
+
# logger.error(f"π₯ Exception type: {type(e).__name__}")
|
181 |
+
|
182 |
+
# # Return error response
|
183 |
+
# return {
|
184 |
+
# "bug_analysis": [{
|
185 |
+
# "line_number": 1,
|
186 |
+
# "error_message": "Analysis failed",
|
187 |
+
# "explanation": f"An error occurred during analysis: {str(e)}",
|
188 |
+
# "fix_suggestion": "Please try again or contact support"
|
189 |
+
# }],
|
190 |
+
# "corrected_code": code,
|
191 |
+
# "error": str(e),
|
192 |
+
# "error_type": type(e).__name__
|
193 |
+
# }
|
194 |
+
|
195 |
+
# analyzer.py
|
196 |
+
|
197 |
+
import torch
|
198 |
+
import json
|
199 |
+
import time
|
200 |
+
import logging
|
201 |
+
|
202 |
+
# Configure logger
|
203 |
+
logger = logging.getLogger("CodeAnalyzer")
|
204 |
+
logger.setLevel(logging.INFO)
|
205 |
+
handler = logging.StreamHandler()
|
206 |
+
formatter = logging.Formatter("[%(asctime)s] [%(levelname)s] - %(message)s")
|
207 |
+
handler.setFormatter(formatter)
|
208 |
+
logger.addHandler(handler)
|
209 |
+
|
210 |
+
|
211 |
+
def analyze_code(tokenizer, model, language, code):
|
212 |
+
start_time = time.time()
|
213 |
+
|
214 |
+
messages = [
|
215 |
+
{
|
216 |
+
"role": "system",
|
217 |
+
"content": (
|
218 |
+
"You are a helpful and expert-level AI code reviewer and bug fixer. "
|
219 |
+
"Your task is to analyze the given buggy code in the specified programming language, "
|
220 |
+
"identify bugs (logical, syntax, runtime, etc.), and fix them. "
|
221 |
+
"Return a JSON object with the following keys:\n\n"
|
222 |
+
"1. 'bug_analysis': a list of objects, each containing:\n"
|
223 |
+
" - 'line_number': the line number (approximate if needed)\n"
|
224 |
+
" - 'error_message': a short name of the bug\n"
|
225 |
+
" - 'explanation': short explanation of the problem\n"
|
226 |
+
" - 'fix_suggestion': how to fix it\n"
|
227 |
+
"2. 'corrected_code': the entire corrected code block.\n\n"
|
228 |
+
"Respond only with a JSON block, no extra commentary."
|
229 |
+
)
|
230 |
+
},
|
231 |
+
{
|
232 |
+
"role": "user",
|
233 |
+
"content": f"π» Language: {language}\nπ Buggy Code:\n```{language.lower()}\n{code.strip()}\n```"
|
234 |
+
}
|
235 |
+
]
|
236 |
+
|
237 |
+
try:
|
238 |
+
logger.info("π¦ Tokenizing input...")
|
239 |
+
inputs = tokenizer.apply_chat_template(
|
240 |
+
messages,
|
241 |
+
add_generation_prompt=True,
|
242 |
+
return_tensors="pt"
|
243 |
+
).to(model.device)
|
244 |
+
|
245 |
+
attention_mask = (inputs != tokenizer.pad_token_id).long()
|
246 |
+
|
247 |
+
logger.info("βοΈ Starting generation...")
|
248 |
+
generation_start = time.time()
|
249 |
+
outputs = model.generate(
|
250 |
+
inputs,
|
251 |
+
attention_mask=attention_mask,
|
252 |
+
max_new_tokens=1024,
|
253 |
+
do_sample=False,
|
254 |
+
pad_token_id=tokenizer.eos_token_id,
|
255 |
+
eos_token_id=tokenizer.eos_token_id
|
256 |
+
)
|
257 |
+
generation_time = time.time() - generation_start
|
258 |
+
logger.info(f"β‘ Generation completed in {generation_time:.2f} seconds")
|
259 |
+
|
260 |
+
logger.info("π Decoding response...")
|
261 |
+
response = tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True)
|
262 |
+
|
263 |
+
logger.info(f"π Response length: {len(response)} characters")
|
264 |
+
logger.info(f"π First 100 chars: {response[:100]}...")
|
265 |
+
|
266 |
+
# Attempt to parse as JSON
|
267 |
+
logger.info("π Attempting to parse JSON...")
|
268 |
+
cleaned_response = response.strip()
|
269 |
+
if cleaned_response.startswith('```json'):
|
270 |
+
cleaned_response = cleaned_response[7:]
|
271 |
+
elif cleaned_response.startswith('```'):
|
272 |
+
cleaned_response = cleaned_response[3:]
|
273 |
+
if cleaned_response.endswith('```'):
|
274 |
+
cleaned_response = cleaned_response[:-3]
|
275 |
+
|
276 |
+
cleaned_response = cleaned_response.strip()
|
277 |
+
|
278 |
+
json_output = json.loads(cleaned_response)
|
279 |
+
|
280 |
+
total_time = time.time() - start_time
|
281 |
+
logger.info(f"β
Analysis completed successfully in {total_time:.2f} seconds")
|
282 |
+
|
283 |
+
# Validate and patch missing keys
|
284 |
+
if not isinstance(json_output, dict):
|
285 |
+
raise ValueError("Parsed response is not a dictionary")
|
286 |
+
|
287 |
+
if 'bug_analysis' not in json_output:
|
288 |
+
logger.warning("β οΈ Missing 'bug_analysis' key, adding empty list")
|
289 |
+
json_output['bug_analysis'] = []
|
290 |
+
|
291 |
+
if 'corrected_code' not in json_output:
|
292 |
+
logger.warning("β οΈ Missing 'corrected_code' key, adding original code")
|
293 |
+
json_output['corrected_code'] = code
|
294 |
+
|
295 |
+
return json_output
|
296 |
+
|
297 |
+
except json.JSONDecodeError as e:
|
298 |
+
logger.error(f"β JSON decode error: {e}")
|
299 |
+
logger.error(f"π Raw response: {repr(response)}")
|
300 |
+
return {
|
301 |
+
"bug_analysis": [{
|
302 |
+
"line_number": 1,
|
303 |
+
"error_message": "Analysis parsing failed",
|
304 |
+
"explanation": "The AI model returned a response that couldn't be parsed as JSON",
|
305 |
+
"fix_suggestion": "Please try again or check the code format"
|
306 |
+
}],
|
307 |
+
"corrected_code": code,
|
308 |
+
"raw_output": response,
|
309 |
+
"parsing_error": str(e)
|
310 |
+
}
|
311 |
+
|
312 |
+
except Exception as e:
|
313 |
+
total_time = time.time() - start_time
|
314 |
+
logger.error(f"β Analysis failed after {total_time:.2f} seconds: {str(e)}")
|
315 |
+
logger.error(f"π₯ Exception type: {type(e).__name__}")
|
316 |
+
return {
|
317 |
+
"bug_analysis": [{
|
318 |
+
"line_number": 1,
|
319 |
+
"error_message": "Analysis failed",
|
320 |
+
"explanation": f"An error occurred during analysis: {str(e)}",
|
321 |
+
"fix_suggestion": "Please try again or contact support"
|
322 |
+
}],
|
323 |
+
"corrected_code": code,
|
324 |
+
"error": str(e),
|
325 |
+
"error_type": type(e).__name__
|
326 |
+
}
|
app.py
ADDED
@@ -0,0 +1,55 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# # app.py
|
2 |
+
|
3 |
+
# from model import load_model
|
4 |
+
# from analyzer import analyze_code
|
5 |
+
# import json
|
6 |
+
|
7 |
+
# if __name__ == "__main__":
|
8 |
+
# print("π§ AI Bug Explainer - Local Terminal Interface")
|
9 |
+
# language = input("Enter programming language (e.g., Python): ")
|
10 |
+
# print("\nPaste your buggy code. End input with a line that says only 'END':\n")
|
11 |
+
|
12 |
+
# lines = []
|
13 |
+
# while True:
|
14 |
+
# line = input()
|
15 |
+
# if line.strip() == "END":
|
16 |
+
# break
|
17 |
+
# lines.append(line)
|
18 |
+
|
19 |
+
# code = "\n".join(lines)
|
20 |
+
|
21 |
+
# tokenizer, model = load_model()
|
22 |
+
# print("\nπ Analyzing your code...\n")
|
23 |
+
# result = analyze_code(language, code, tokenizer, model)
|
24 |
+
|
25 |
+
# print(json.dumps(result, indent=2))
|
26 |
+
# app.py
|
27 |
+
|
28 |
+
from model import load_model
|
29 |
+
from analyzer import analyze_code
|
30 |
+
import json
|
31 |
+
|
32 |
+
def main():
|
33 |
+
print("π§ Loading model...")
|
34 |
+
tokenizer, model = load_model()
|
35 |
+
|
36 |
+
print("\nπ₯ Enter your code for analysis.")
|
37 |
+
language = input("Programming Language (e.g., Python, JavaScript): ").strip()
|
38 |
+
|
39 |
+
print("Paste your buggy code (end input with an empty line):")
|
40 |
+
code_lines = []
|
41 |
+
while True:
|
42 |
+
line = input()
|
43 |
+
if line == "":
|
44 |
+
break
|
45 |
+
code_lines.append(line)
|
46 |
+
code = "\n".join(code_lines)
|
47 |
+
|
48 |
+
print("\nπ Analyzing your code...\n")
|
49 |
+
result = analyze_code(tokenizer, model, language, code)
|
50 |
+
|
51 |
+
print("\nπ§Ύ JSON Response:")
|
52 |
+
print(json.dumps(result, indent=2))
|
53 |
+
|
54 |
+
if __name__ == "__main__":
|
55 |
+
main()
|
main.py
ADDED
@@ -0,0 +1,409 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# from fastapi import FastAPI, HTTPException
|
2 |
+
# from fastapi.middleware.cors import CORSMiddleware
|
3 |
+
# from pydantic import BaseModel
|
4 |
+
# from model import load_model
|
5 |
+
# from analyzer import analyze_code
|
6 |
+
# import logging
|
7 |
+
|
8 |
+
# app = FastAPI(
|
9 |
+
# title="AI Bug Explainer",
|
10 |
+
# description="An AI service that detects and fixes bugs in code",
|
11 |
+
# version="1.0.0"
|
12 |
+
# )
|
13 |
+
|
14 |
+
# # CORS setup
|
15 |
+
# app.add_middleware(
|
16 |
+
# CORSMiddleware,
|
17 |
+
# allow_origins=["*"], # Replace with your frontend URL in prod
|
18 |
+
# allow_credentials=True,
|
19 |
+
# allow_methods=["*"],
|
20 |
+
# allow_headers=["*"],
|
21 |
+
# )
|
22 |
+
|
23 |
+
# # Logging setup
|
24 |
+
# logging.basicConfig(level=logging.INFO)
|
25 |
+
|
26 |
+
# class AnalyzeRequest(BaseModel):
|
27 |
+
# language: str
|
28 |
+
# code: str
|
29 |
+
|
30 |
+
# @app.post("/analyze")
|
31 |
+
# async def analyze(req: AnalyzeRequest):
|
32 |
+
# logging.info(f"π Received code for analysis ({req.language})")
|
33 |
+
|
34 |
+
# result = analyze_code(req.language, req.code, tokenizer, model)
|
35 |
+
|
36 |
+
# if result is None:
|
37 |
+
# raise HTTPException(status_code=500, detail="Model failed to return any response.")
|
38 |
+
|
39 |
+
# if not isinstance(result, dict):
|
40 |
+
# logging.warning("β οΈ Model did not return valid JSON, sending raw output")
|
41 |
+
# return {
|
42 |
+
# "bugs": [],
|
43 |
+
# "corrected_code": "",
|
44 |
+
# "raw_output": result
|
45 |
+
# }
|
46 |
+
|
47 |
+
# return {
|
48 |
+
# "bugs": result.get("bug_analysis", []),
|
49 |
+
# "corrected_code": result.get("corrected_code", ""),
|
50 |
+
# "raw_output": "" # So frontend doesn't break
|
51 |
+
# }
|
52 |
+
|
53 |
+
# # Load model
|
54 |
+
# print("π§ Loading model...")
|
55 |
+
# tokenizer, model = load_model()
|
56 |
+
# print("β
Model loaded!")
|
57 |
+
|
58 |
+
# from fastapi import FastAPI, HTTPException
|
59 |
+
# from fastapi.middleware.cors import CORSMiddleware
|
60 |
+
# from pydantic import BaseModel
|
61 |
+
# from model import load_model
|
62 |
+
# from analyzer import analyze_code
|
63 |
+
# import logging
|
64 |
+
|
65 |
+
# app = FastAPI(
|
66 |
+
# title="AI Bug Explainer ML Microservice",
|
67 |
+
# description="An AI service that detects and fixes bugs in code",
|
68 |
+
# version="1.0.0"
|
69 |
+
# )
|
70 |
+
|
71 |
+
# # CORS setup
|
72 |
+
# app.add_middleware(
|
73 |
+
# CORSMiddleware,
|
74 |
+
# allow_origins=["*"], # Replace with your frontend URL in prod
|
75 |
+
# allow_credentials=True,
|
76 |
+
# allow_methods=["*"],
|
77 |
+
# allow_headers=["*"],
|
78 |
+
# )
|
79 |
+
|
80 |
+
# # Logging setup
|
81 |
+
# logging.basicConfig(level=logging.INFO)
|
82 |
+
|
83 |
+
# class AnalyzeRequest(BaseModel):
|
84 |
+
# language: str
|
85 |
+
# code: str
|
86 |
+
|
87 |
+
# # Transform bug analysis to match frontend expectations
|
88 |
+
# def transform_bug_to_issue(bug):
|
89 |
+
# """Transform ML service bug format to frontend issue format"""
|
90 |
+
# return {
|
91 |
+
# "lineNumber": bug.get("line_number", 0),
|
92 |
+
# "type": bug.get("error_message", "Unknown Error"),
|
93 |
+
# "message": bug.get("explanation", "No explanation provided"),
|
94 |
+
# "suggestion": bug.get("fix_suggestion", "No suggestion provided")
|
95 |
+
# }
|
96 |
+
|
97 |
+
# # Keep your original endpoint for backward compatibility
|
98 |
+
# @app.post("/analyze")
|
99 |
+
# async def analyze(req: AnalyzeRequest):
|
100 |
+
# logging.info(f"π Received code for analysis ({req.language})")
|
101 |
+
|
102 |
+
# result = analyze_code(req.language, req.code, tokenizer, model)
|
103 |
+
|
104 |
+
# if result is None:
|
105 |
+
# raise HTTPException(status_code=500, detail="Model failed to return any response.")
|
106 |
+
|
107 |
+
# if not isinstance(result, dict):
|
108 |
+
# logging.warning("β οΈ Model did not return valid JSON, sending raw output")
|
109 |
+
# return {
|
110 |
+
# "bugs": [],
|
111 |
+
# "corrected_code": "",
|
112 |
+
# "raw_output": result
|
113 |
+
# }
|
114 |
+
|
115 |
+
# return {
|
116 |
+
# "bugs": result.get("bug_analysis", []),
|
117 |
+
# "corrected_code": result.get("corrected_code", ""),
|
118 |
+
# "raw_output": "" # So frontend doesn't break
|
119 |
+
# }
|
120 |
+
|
121 |
+
# # NEW: Add frontend-compatible endpoint
|
122 |
+
# @app.post("/analysis/submit")
|
123 |
+
# async def analyze_for_frontend(req: AnalyzeRequest):
|
124 |
+
# logging.info(f"π Frontend: Received code for analysis ({req.language})")
|
125 |
+
|
126 |
+
# result = analyze_code(req.language, req.code, tokenizer, model)
|
127 |
+
|
128 |
+
# if result is None:
|
129 |
+
# raise HTTPException(status_code=500, detail="Model failed to return any response.")
|
130 |
+
|
131 |
+
# # If result is not valid JSON, return raw output as fallback
|
132 |
+
# if not isinstance(result, dict):
|
133 |
+
# logging.warning("β οΈ Model did not return valid JSON, showing raw output")
|
134 |
+
# return {
|
135 |
+
# "success": False,
|
136 |
+
# "has_json_output": False,
|
137 |
+
# "corrected_code": "",
|
138 |
+
# "issues": [],
|
139 |
+
# "raw_output": str(result)
|
140 |
+
# }
|
141 |
+
|
142 |
+
# # Successfully parsed JSON
|
143 |
+
# bugs = result.get("bug_analysis", [])
|
144 |
+
# issues = [transform_bug_to_issue(bug) for bug in bugs]
|
145 |
+
# corrected_code = result.get("corrected_code", "")
|
146 |
+
|
147 |
+
# return {
|
148 |
+
# "success": True,
|
149 |
+
# "has_json_output": True,
|
150 |
+
# "corrected_code": corrected_code,
|
151 |
+
# "issues": issues,
|
152 |
+
# "raw_output": ""
|
153 |
+
# }
|
154 |
+
|
155 |
+
# # Add history endpoint (placeholder for now)
|
156 |
+
# @app.get("/analysis/history")
|
157 |
+
# async def get_analysis_history():
|
158 |
+
# # TODO: Implement database storage for history
|
159 |
+
# # For now, return empty array to match frontend expectations
|
160 |
+
# return {"data": []}
|
161 |
+
|
162 |
+
# # Health check endpoint
|
163 |
+
# @app.get("/health")
|
164 |
+
# async def health_check():
|
165 |
+
# return {
|
166 |
+
# "status": "healthy",
|
167 |
+
# "model_loaded": tokenizer is not None and model is not None
|
168 |
+
# }
|
169 |
+
|
170 |
+
# # Load model
|
171 |
+
# print("π§ Loading model...")
|
172 |
+
# tokenizer, model = load_model()
|
173 |
+
# print("β
Model loaded!")
|
174 |
+
|
175 |
+
# if __name__ == "__main__":
|
176 |
+
# import uvicorn
|
177 |
+
# uvicorn.run(app, host="0.0.0.0", port=8000)
|
178 |
+
|
179 |
+
from fastapi import FastAPI, HTTPException, BackgroundTasks
|
180 |
+
from fastapi.middleware.cors import CORSMiddleware
|
181 |
+
from pydantic import BaseModel
|
182 |
+
from model import load_model_async, get_model, is_model_loaded, get_model_info
|
183 |
+
from analyzer import analyze_code
|
184 |
+
import logging
|
185 |
+
import asyncio
|
186 |
+
import time
|
187 |
+
from dotenv import load_dotenv
|
188 |
+
load_dotenv()
|
189 |
+
|
190 |
+
# Configure logging
|
191 |
+
logging.basicConfig(
|
192 |
+
level=logging.INFO,
|
193 |
+
format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
|
194 |
+
)
|
195 |
+
logger = logging.getLogger(__name__)
|
196 |
+
|
197 |
+
app = FastAPI(
|
198 |
+
title="AI Bug Explainer ML Microservice",
|
199 |
+
description="An AI service that detects and fixes bugs in code",
|
200 |
+
version="1.0.0"
|
201 |
+
)
|
202 |
+
|
203 |
+
# CORS setup
|
204 |
+
app.add_middleware(
|
205 |
+
CORSMiddleware,
|
206 |
+
allow_origins=["*"], # Replace with your frontend URL in prod
|
207 |
+
allow_credentials=True,
|
208 |
+
allow_methods=["*"],
|
209 |
+
allow_headers=["*"],
|
210 |
+
)
|
211 |
+
|
212 |
+
class AnalyzeRequest(BaseModel):
|
213 |
+
language: str
|
214 |
+
code: str
|
215 |
+
|
216 |
+
# Global variables for model loading status
|
217 |
+
model_load_start_time = None
|
218 |
+
model_load_task = None
|
219 |
+
|
220 |
+
def transform_bug_to_issue(bug):
|
221 |
+
"""Transform ML service bug format to frontend issue format"""
|
222 |
+
return {
|
223 |
+
"lineNumber": bug.get("line_number", 0),
|
224 |
+
"type": bug.get("error_message", "Unknown Error"),
|
225 |
+
"message": bug.get("explanation", "No explanation provided"),
|
226 |
+
"suggestion": bug.get("fix_suggestion", "No suggestion provided")
|
227 |
+
}
|
228 |
+
|
229 |
+
@app.on_event("startup")
|
230 |
+
async def startup_event():
|
231 |
+
"""Start model loading in background when server starts"""
|
232 |
+
global model_load_start_time, model_load_task
|
233 |
+
logger.info("π Starting ML microservice...")
|
234 |
+
logger.info("π§ Initiating background model loading...")
|
235 |
+
|
236 |
+
model_load_start_time = time.time()
|
237 |
+
|
238 |
+
# Start model loading in background
|
239 |
+
model_load_task = asyncio.create_task(load_model_async())
|
240 |
+
|
241 |
+
logger.info("β
Server started! Model is loading in background...")
|
242 |
+
|
243 |
+
@app.get("/health")
|
244 |
+
async def health_check():
|
245 |
+
"""Enhanced health check with model loading status"""
|
246 |
+
global model_load_start_time
|
247 |
+
|
248 |
+
model_info = get_model_info()
|
249 |
+
loading_time = None
|
250 |
+
|
251 |
+
if model_load_start_time:
|
252 |
+
loading_time = round(time.time() - model_load_start_time, 2)
|
253 |
+
|
254 |
+
return {
|
255 |
+
"status": "healthy",
|
256 |
+
"model_info": model_info,
|
257 |
+
"loading_time_seconds": loading_time,
|
258 |
+
"ready_for_inference": model_info["loaded"]
|
259 |
+
}
|
260 |
+
|
261 |
+
@app.get("/model/status")
|
262 |
+
async def model_status():
|
263 |
+
"""Get detailed model loading status"""
|
264 |
+
global model_load_start_time
|
265 |
+
|
266 |
+
model_info = get_model_info()
|
267 |
+
loading_time = None
|
268 |
+
|
269 |
+
if model_load_start_time:
|
270 |
+
loading_time = round(time.time() - model_load_start_time, 2)
|
271 |
+
|
272 |
+
return {
|
273 |
+
"model_id": model_info["model_id"],
|
274 |
+
"loaded": model_info["loaded"],
|
275 |
+
"loading": model_info["loading"],
|
276 |
+
"loading_time_seconds": loading_time,
|
277 |
+
"ready": model_info["loaded"]
|
278 |
+
}
|
279 |
+
|
280 |
+
@app.post("/analyze")
|
281 |
+
async def analyze(req: AnalyzeRequest):
|
282 |
+
"""Original analyze endpoint with model loading check"""
|
283 |
+
logger.info(f"π Received code for analysis ({req.language})")
|
284 |
+
|
285 |
+
# Check if model is loaded
|
286 |
+
if not is_model_loaded():
|
287 |
+
# Wait for model to load (with timeout)
|
288 |
+
try:
|
289 |
+
await asyncio.wait_for(model_load_task, timeout=300) # 5 minute timeout
|
290 |
+
except asyncio.TimeoutError:
|
291 |
+
raise HTTPException(
|
292 |
+
status_code=503,
|
293 |
+
detail="Model is still loading. Please try again in a few moments."
|
294 |
+
)
|
295 |
+
|
296 |
+
try:
|
297 |
+
tokenizer, model = get_model()
|
298 |
+
result = analyze_code(req.language, req.code, tokenizer, model)
|
299 |
+
|
300 |
+
if result is None:
|
301 |
+
raise HTTPException(status_code=500, detail="Model failed to return any response.")
|
302 |
+
|
303 |
+
if not isinstance(result, dict):
|
304 |
+
logger.warning("β οΈ Model did not return valid JSON, sending raw output")
|
305 |
+
return {
|
306 |
+
"bugs": [],
|
307 |
+
"corrected_code": "",
|
308 |
+
"raw_output": result
|
309 |
+
}
|
310 |
+
|
311 |
+
return {
|
312 |
+
"bugs": result.get("bug_analysis", []),
|
313 |
+
"corrected_code": result.get("corrected_code", ""),
|
314 |
+
"raw_output": ""
|
315 |
+
}
|
316 |
+
except Exception as e:
|
317 |
+
logger.error(f"Analysis error: {e}")
|
318 |
+
raise HTTPException(status_code=500, detail=f"Analysis failed: {str(e)}")
|
319 |
+
|
320 |
+
@app.post("/analysis/submit")
|
321 |
+
async def analyze_for_frontend(req: AnalyzeRequest):
|
322 |
+
"""Frontend-compatible endpoint with model loading check"""
|
323 |
+
logger.info(f"π Frontend: Received code for analysis ({req.language})")
|
324 |
+
|
325 |
+
# Check if model is loaded
|
326 |
+
if not is_model_loaded():
|
327 |
+
# If model is still loading, return appropriate response
|
328 |
+
if model_load_task and not model_load_task.done():
|
329 |
+
return {
|
330 |
+
"success": False,
|
331 |
+
"has_json_output": False,
|
332 |
+
"corrected_code": "",
|
333 |
+
"issues": [],
|
334 |
+
"raw_output": "Model is still loading. Please wait a moment and try again.",
|
335 |
+
"model_status": "loading"
|
336 |
+
}
|
337 |
+
else:
|
338 |
+
# Try to wait for model loading
|
339 |
+
try:
|
340 |
+
await asyncio.wait_for(model_load_task, timeout=30) # Short timeout for frontend
|
341 |
+
except (asyncio.TimeoutError, Exception):
|
342 |
+
return {
|
343 |
+
"success": False,
|
344 |
+
"has_json_output": False,
|
345 |
+
"corrected_code": "",
|
346 |
+
"issues": [],
|
347 |
+
"raw_output": "Model is not ready yet. Please try again in a few moments.",
|
348 |
+
"model_status": "loading"
|
349 |
+
}
|
350 |
+
|
351 |
+
try:
|
352 |
+
tokenizer, model = get_model()
|
353 |
+
result = analyze_code(req.language, req.code, tokenizer, model)
|
354 |
+
|
355 |
+
if result is None:
|
356 |
+
return {
|
357 |
+
"success": False,
|
358 |
+
"has_json_output": False,
|
359 |
+
"corrected_code": "",
|
360 |
+
"issues": [],
|
361 |
+
"raw_output": "Model failed to return any response.",
|
362 |
+
"model_status": "error"
|
363 |
+
}
|
364 |
+
|
365 |
+
# If result is not valid JSON, return raw output as fallback
|
366 |
+
if not isinstance(result, dict):
|
367 |
+
logger.warning("β οΈ Model did not return valid JSON, showing raw output")
|
368 |
+
return {
|
369 |
+
"success": False,
|
370 |
+
"has_json_output": False,
|
371 |
+
"corrected_code": "",
|
372 |
+
"issues": [],
|
373 |
+
"raw_output": str(result),
|
374 |
+
"model_status": "loaded"
|
375 |
+
}
|
376 |
+
|
377 |
+
# Successfully parsed JSON
|
378 |
+
bugs = result.get("bug_analysis", [])
|
379 |
+
issues = [transform_bug_to_issue(bug) for bug in bugs]
|
380 |
+
corrected_code = result.get("corrected_code", "")
|
381 |
+
|
382 |
+
return {
|
383 |
+
"success": True,
|
384 |
+
"has_json_output": True,
|
385 |
+
"corrected_code": corrected_code,
|
386 |
+
"issues": issues,
|
387 |
+
"raw_output": "",
|
388 |
+
"model_status": "loaded"
|
389 |
+
}
|
390 |
+
|
391 |
+
except Exception as e:
|
392 |
+
logger.error(f"Frontend analysis error: {e}")
|
393 |
+
return {
|
394 |
+
"success": False,
|
395 |
+
"has_json_output": False,
|
396 |
+
"corrected_code": "",
|
397 |
+
"issues": [],
|
398 |
+
"raw_output": f"Analysis failed: {str(e)}",
|
399 |
+
"model_status": "error"
|
400 |
+
}
|
401 |
+
|
402 |
+
@app.get("/analysis/history")
|
403 |
+
async def get_analysis_history():
|
404 |
+
"""Get analysis history (placeholder)"""
|
405 |
+
return {"data": []}
|
406 |
+
|
407 |
+
if __name__ == "__main__":
|
408 |
+
import uvicorn
|
409 |
+
uvicorn.run(app, host="0.0.0.0", port=8000)
|
model.py
ADDED
@@ -0,0 +1,124 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# model.py - Optimized version
|
2 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
3 |
+
import torch
|
4 |
+
from functools import lru_cache
|
5 |
+
import os
|
6 |
+
import asyncio
|
7 |
+
from concurrent.futures import ThreadPoolExecutor
|
8 |
+
import logging
|
9 |
+
|
10 |
+
logger = logging.getLogger(__name__)
|
11 |
+
|
12 |
+
# Global variables to store loaded model
|
13 |
+
_tokenizer = None
|
14 |
+
_model = None
|
15 |
+
_model_loading = False
|
16 |
+
_model_loaded = False
|
17 |
+
|
18 |
+
@lru_cache(maxsize=1)
|
19 |
+
def get_model_config():
|
20 |
+
"""Cache model configuration"""
|
21 |
+
return {
|
22 |
+
"model_id": "deepseek-ai/deepseek-coder-1.3b-instruct",
|
23 |
+
"torch_dtype": torch.bfloat16,
|
24 |
+
"device_map": "auto",
|
25 |
+
"trust_remote_code": True,
|
26 |
+
# Add these optimizations
|
27 |
+
"low_cpu_mem_usage": True,
|
28 |
+
"use_cache": True,
|
29 |
+
}
|
30 |
+
|
31 |
+
def load_model_sync():
|
32 |
+
"""Synchronous model loading with optimizations"""
|
33 |
+
global _tokenizer, _model, _model_loaded
|
34 |
+
|
35 |
+
if _model_loaded:
|
36 |
+
return _tokenizer, _model
|
37 |
+
|
38 |
+
config = get_model_config()
|
39 |
+
model_id = config["model_id"]
|
40 |
+
|
41 |
+
logger.info(f"π§ Loading model {model_id}...")
|
42 |
+
|
43 |
+
try:
|
44 |
+
# Set cache directory to avoid re-downloading
|
45 |
+
cache_dir = os.environ.get("TRANSFORMERS_CACHE", "./model_cache")
|
46 |
+
os.makedirs(cache_dir, exist_ok=True)
|
47 |
+
|
48 |
+
# Load tokenizer first (faster)
|
49 |
+
logger.info("π Loading tokenizer...")
|
50 |
+
_tokenizer = AutoTokenizer.from_pretrained(
|
51 |
+
model_id,
|
52 |
+
trust_remote_code=config["trust_remote_code"],
|
53 |
+
cache_dir=cache_dir,
|
54 |
+
use_fast=True, # Use fast tokenizer if available
|
55 |
+
)
|
56 |
+
|
57 |
+
# Load model with optimizations
|
58 |
+
logger.info("π§ Loading model...")
|
59 |
+
_model = AutoModelForCausalLM.from_pretrained(
|
60 |
+
model_id,
|
61 |
+
trust_remote_code=config["trust_remote_code"],
|
62 |
+
torch_dtype=config["torch_dtype"],
|
63 |
+
device_map=config["device_map"],
|
64 |
+
low_cpu_mem_usage=config["low_cpu_mem_usage"],
|
65 |
+
cache_dir=cache_dir,
|
66 |
+
offload_folder="offload",
|
67 |
+
offload_state_dict=True
|
68 |
+
)
|
69 |
+
|
70 |
+
# Set to evaluation mode
|
71 |
+
_model.eval()
|
72 |
+
|
73 |
+
_model_loaded = True
|
74 |
+
logger.info("β
Model loaded successfully!")
|
75 |
+
return _tokenizer, _model
|
76 |
+
|
77 |
+
except Exception as e:
|
78 |
+
logger.error(f"β Failed to load model: {e}")
|
79 |
+
raise
|
80 |
+
|
81 |
+
async def load_model_async():
|
82 |
+
"""Asynchronous model loading"""
|
83 |
+
global _model_loading
|
84 |
+
|
85 |
+
if _model_loaded:
|
86 |
+
return _tokenizer, _model
|
87 |
+
|
88 |
+
if _model_loading:
|
89 |
+
# Wait for ongoing loading to complete
|
90 |
+
while _model_loading and not _model_loaded:
|
91 |
+
await asyncio.sleep(0.1)
|
92 |
+
return _tokenizer, _model
|
93 |
+
|
94 |
+
_model_loading = True
|
95 |
+
|
96 |
+
try:
|
97 |
+
# Run model loading in thread pool to avoid blocking
|
98 |
+
loop = asyncio.get_event_loop()
|
99 |
+
with ThreadPoolExecutor(max_workers=1) as executor:
|
100 |
+
tokenizer, model = await loop.run_in_executor(
|
101 |
+
executor, load_model_sync
|
102 |
+
)
|
103 |
+
return tokenizer, model
|
104 |
+
finally:
|
105 |
+
_model_loading = False
|
106 |
+
|
107 |
+
def get_model():
|
108 |
+
"""Get the loaded model (for synchronous access)"""
|
109 |
+
if not _model_loaded:
|
110 |
+
return load_model_sync()
|
111 |
+
return _tokenizer, _model
|
112 |
+
|
113 |
+
def is_model_loaded():
|
114 |
+
"""Check if model is loaded"""
|
115 |
+
return _model_loaded
|
116 |
+
|
117 |
+
def get_model_info():
|
118 |
+
"""Get model information without loading"""
|
119 |
+
config = get_model_config()
|
120 |
+
return {
|
121 |
+
"model_id": config["model_id"],
|
122 |
+
"loaded": _model_loaded,
|
123 |
+
"loading": _model_loading,
|
124 |
+
}
|
requirements.txt
ADDED
@@ -0,0 +1,24 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# torch>=2.1.0
|
2 |
+
# transformers>=4.40.0
|
3 |
+
# accelerate>=0.25.0
|
4 |
+
# bitsandbytes
|
5 |
+
# fastapi
|
6 |
+
# uvicorn
|
7 |
+
#Your original dependencies (optimized versions)
|
8 |
+
torch>=2.1.0
|
9 |
+
transformers==4.41.1
|
10 |
+
accelerate==0.30.1
|
11 |
+
bitsandbytes
|
12 |
+
fastapi
|
13 |
+
uvicorn[standard]
|
14 |
+
|
15 |
+
# Additional optimizations for faster loading
|
16 |
+
tokenizers>=0.15.0 # Fast tokenizers (auto-installed with transformers but explicit for optimization)
|
17 |
+
safetensors>=0.4.0 # Faster model loading format
|
18 |
+
huggingface-hub>=0.19.0 # Better caching and download management
|
19 |
+
|
20 |
+
# Optional performance improvements
|
21 |
+
psutil>=5.9.0 # For system monitoring
|
22 |
+
python-multipart # For FastAPI file uploads if needed
|
23 |
+
|
24 |
+
python-dotenv
|
setup.py
ADDED
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python3
|
2 |
+
"""
|
3 |
+
Quick setup script to optimize your existing ML microservice.
|
4 |
+
Run this to set up caching and pre-download the model.
|
5 |
+
"""
|
6 |
+
|
7 |
+
import os
|
8 |
+
import sys
|
9 |
+
import logging
|
10 |
+
from pathlib import Path
|
11 |
+
|
12 |
+
# Configure logging
|
13 |
+
logging.basicConfig(level=logging.INFO)
|
14 |
+
logger = logging.getLogger(__name__)
|
15 |
+
|
16 |
+
def setup_cache_directory():
|
17 |
+
"""Create cache directory for models"""
|
18 |
+
cache_dir = Path("./model_cache")
|
19 |
+
cache_dir.mkdir(exist_ok=True)
|
20 |
+
logger.info(f"β
Cache directory created: {cache_dir.absolute()}")
|
21 |
+
return cache_dir
|
22 |
+
|
23 |
+
def set_environment_variables():
|
24 |
+
"""Set environment variables for optimization"""
|
25 |
+
env_vars = {
|
26 |
+
"TRANSFORMERS_CACHE": "./model_cache",
|
27 |
+
"HF_HOME": "./model_cache",
|
28 |
+
"TORCH_HOME": "./model_cache",
|
29 |
+
"TOKENIZERS_PARALLELISM": "false",
|
30 |
+
"OMP_NUM_THREADS": "4"
|
31 |
+
}
|
32 |
+
|
33 |
+
for key, value in env_vars.items():
|
34 |
+
os.environ[key] = value
|
35 |
+
logger.info(f"Set {key}={value}")
|
36 |
+
|
37 |
+
def pre_download_model():
|
38 |
+
"""Pre-download the model to cache"""
|
39 |
+
try:
|
40 |
+
from transformers import AutoTokenizer, AutoModelForCausalLM
|
41 |
+
|
42 |
+
model_id = "deepseek-ai/deepseek-coder-1.3b-instruct"
|
43 |
+
cache_dir = "./model_cache"
|
44 |
+
|
45 |
+
logger.info(f"π§ Pre-downloading model: {model_id}")
|
46 |
+
logger.info("This may take a few minutes on first run...")
|
47 |
+
|
48 |
+
# Download tokenizer
|
49 |
+
logger.info("π Downloading tokenizer...")
|
50 |
+
tokenizer = AutoTokenizer.from_pretrained(
|
51 |
+
model_id,
|
52 |
+
cache_dir=cache_dir,
|
53 |
+
trust_remote_code=True
|
54 |
+
)
|
55 |
+
|
56 |
+
# Download model
|
57 |
+
logger.info("π§ Downloading model...")
|
58 |
+
model = AutoModelForCausalLM.from_pretrained(
|
59 |
+
model_id,
|
60 |
+
cache_dir=cache_dir,
|
61 |
+
trust_remote_code=True,
|
62 |
+
torch_dtype="auto", # Let it choose the best dtype
|
63 |
+
low_cpu_mem_usage=True,
|
64 |
+
)
|
65 |
+
|
66 |
+
logger.info("β
Model downloaded and cached successfully!")
|
67 |
+
logger.info(f"π Model cached in: {Path(cache_dir).absolute()}")
|
68 |
+
|
69 |
+
# Test that everything works
|
70 |
+
logger.info("π§ͺ Testing model loading...")
|
71 |
+
del model, tokenizer # Free memory
|
72 |
+
|
73 |
+
return True
|
74 |
+
|
75 |
+
except Exception as e:
|
76 |
+
logger.error(f"β Failed to pre-download model: {e}")
|
77 |
+
return False
|
78 |
+
|
79 |
+
def main():
|
80 |
+
"""Main setup function"""
|
81 |
+
logger.info("π Setting up ML Microservice Optimizations")
|
82 |
+
logger.info("=" * 50)
|
83 |
+
|
84 |
+
# Step 1: Setup cache directory
|
85 |
+
setup_cache_directory()
|
86 |
+
|
87 |
+
# Step 2: Set environment variables
|
88 |
+
set_environment_variables()
|
89 |
+
|
90 |
+
# Step 3: Pre-download model
|
91 |
+
success = pre_download_model()
|
92 |
+
|
93 |
+
if success:
|
94 |
+
logger.info("\nβ
Setup completed successfully!")
|
95 |
+
logger.info("π Next steps:")
|
96 |
+
logger.info("1. Replace your main.py with the optimized version")
|
97 |
+
logger.info("2. Replace your model.py with the optimized version")
|
98 |
+
logger.info("3. Run: python main.py")
|
99 |
+
logger.info("\nπ Your server will now start much faster!")
|
100 |
+
else:
|
101 |
+
logger.error("\nβ Setup failed!")
|
102 |
+
logger.error("Please check your internet connection and try again.")
|
103 |
+
sys.exit(1)
|
104 |
+
|
105 |
+
if __name__ == "__main__":
|
106 |
+
main()
|