Sushwetabm commited on
Commit
f59cf24
Β·
0 Parent(s):

Deploy ML microservice to Hugging Face Space

Browse files
Files changed (10) hide show
  1. .dockerignore +12 -0
  2. .gitignore +9 -0
  3. Dockerfile +32 -0
  4. __init__.py +0 -0
  5. analyzer.py +326 -0
  6. app.py +55 -0
  7. main.py +409 -0
  8. model.py +124 -0
  9. requirements.txt +24 -0
  10. 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()