File size: 4,517 Bytes
0d6bb64 04e0693 0d6bb64 938dad5 0d6bb64 938dad5 0d6bb64 938dad5 0d6bb64 938dad5 |
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 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
import time
# --- MODEL CONSTANTS ---
MODEL_NAME = "deepseek-ai/deepseek-coder-1.3b-instruct"
# CRITICAL: Force model to use CPU for the free tier
DEVICE = "cpu"
MAX_NEW_TOKENS = 512 # Limit output size for speed and cost control
TORCH_DTYPE = torch.float32 # Use standard float for maximum CPU compatibility
# Global variables for model and tokenizer
model = None
tokenizer = None
# --- API Data Structure ---
class CodeRequest(BaseModel):
"""Defines the expected input structure from the front-end website."""
user_prompt: str # The user's request (e.g., "Fix the bug in this function")
code_context: str # The block of code the user provided
# --- FastAPI App Setup ---
app = FastAPI(title="CodeFlow AI Agent Backend - DeepSeek SLM")
@app.on_event("startup")
async def startup_event():
"""Load the DeepSeek SLM Model and Tokenizer ONLY ONCE when the server starts."""
global model, tokenizer
print(f"--- Starting CodeFlow AI Agent (DeepSeek 1.3B) on {DEVICE} ---")
start_time = time.time()
try:
# Load the Tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
# Load the Model
# Using device_map="cpu" is essential for the free tier.
model = AutoModelForCausalLM.from_pretrained(
MODEL_NAME,
dtype=TORCH_DTYPE, # Changed the argument name from torch_dtype to dtype
device_map=DEVICE,
trust_remote_code=True
)
model.eval() # Set model to evaluation mode
print(f"DeepSeek Model loaded successfully in {time.time() - start_time:.2f} seconds.")
except Exception as e:
# If the model fails to load, log the error and stop
print(f"ERROR: Failed to load DeepSeek model on CPU: {e}")
# Raising an exception will cause the Docker container to exit, which is the correct behavior
# when a critical component (the model) fails to load.
raise RuntimeError(f"Model failed to load on startup: {e}")
# --- CRITICAL HEALTH CHECK ENDPOINT (Fixes the Launch Timeout Error) ---
# This must be outside of any other function and use correct indentation.
@app.get("/")
def read_root():
"""Simple health check endpoint for Hugging Face to confirm the app is running."""
return {"status": "ok", "agent": "CodeBuddy DeepSeek-Coder-1.3B"}
# --- The Code Fixing API Endpoint ---
@app.post("/fix_code")
async def fix_code_endpoint(request: CodeRequest):
"""
Accepts code context and task, processes it with DeepSeek-Coder, and returns the fix.
"""
if model is None or tokenizer is None:
raise HTTPException(status_code=503, detail="AI Agent is still loading or failed to start.")
# --- CONSTRUCT AGENT PROMPT (DeepSeek Instruction Format) ---
instruction = (
f"You are Arya's CodeBuddy, an elite Full-Stack Software Engineer. Your only job is to analyze "
f"the user's request and provide the complete, fixed, or generated code. You must ONLY output "
f"a single, complete, and corrected Markdown code block. Use a friendly and encouraging tone.\n\n"
f"TASK: {request.user_prompt}\n\n"
f"CODE_CONTEXT:\n{request.code_context}"
)
# Format the prompt correctly for the model
prompt = f"### Instruction:\n{instruction}\n\n### Response:\n"
# Tokenize and send tensors to CPU
model_inputs = tokenizer([prompt], return_tensors="pt").to(DEVICE)
try:
# --- GENERATE CODE (CPU Inference) ---
generated_ids = model.generate(
**model_inputs,
max_new_tokens=MAX_NEW_TOKENS,
do_sample=False, # Deterministic output
temperature=0.1, # Low temperature for reliable coding
)
# Decode the output
response_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
# Post-processing: Extract ONLY the text after the '### Response:' tag.
final_code_only = response_text.split("### Response:")[1].strip()
return {"fixed_code": final_code_only}
except Exception as e:
print(f"Generation error: {e}")
# Return a generic error to the user
raise HTTPException(status_code=500, detail="The DeepSeek CodeBuddy encountered a processing error.") |