Update main.py

main.py

@@ -6,17 +6,33 @@ import autopep8
 import subprocess
 import time
 import re
+import os
+from pathlib import Path
 
 app = FastAPI(title="Code Evaluation & Optimization API")
 
+# --- Environment Setup ---
+CACHE_DIR = Path("/.cache/huggingface")
+CACHE_DIR.mkdir(parents=True, exist_ok=True)
+os.environ["TRANSFORMERS_CACHE"] = str(CACHE_DIR)
+os.environ["HF_HOME"] = str(CACHE_DIR)
+
 # --- Load AI Model ---
 MODEL_NAME = "codellama/CodeLlama-7b-hf"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, 
-    device_map="auto", 
-    torch_dtype=torch.float16
-)
+
+try:
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        cache_dir=str(CACHE_DIR)
+    )
+    model = AutoModelForCausalLM.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",
+        torch_dtype=torch.float16,
+        cache_dir=str(CACHE_DIR)
+    )
+except Exception as e:
+    raise RuntimeError(f"Failed to load model: {str(e)}")
 
 # --- Request Models ---
 class CodeRequest(BaseModel):
@@ -105,27 +121,30 @@ def evaluate_code(user_code: str, lang: str) -> dict:
 
 def optimize_code_ai(user_code: str, lang: str) -> str:
     """Generate optimized code using AI and formatting"""
-    # Basic formatting first
-    if lang == "python":
-        user_code = autopep8.fix_code(user_code)
-        user_code = re.sub(r"eval\((.*)\)", r"int(\1)  # Removed eval for security", user_code)
-        user_code = re.sub(r"/ 0", "/ 1  # Fixed division by zero", user_code)
-    
-    # AI-powered optimization
-    prompt = f"Optimize this {lang} code for efficiency and security:\n```{lang}\n{user_code}\n```\nOptimized version:"
-    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
-    
-    with torch.no_grad():
-        outputs = model.generate(**inputs, max_length=1024)
-    
-    optimized_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    
-    # Extract just the code block if LLM added explanation
-    code_match = re.search(r'```(?:python)?\n(.*?)\n```', optimized_code, re.DOTALL)
-    if code_match:
-        optimized_code = code_match.group(1)
-    
-    return optimized_code if optimized_code else user_code
+    try:
+        # Basic formatting first
+        if lang == "python":
+            user_code = autopep8.fix_code(user_code)
+            user_code = re.sub(r"eval\((.*)\)", r"int(\1)  # Removed eval for security", user_code)
+            user_code = re.sub(r"/ 0", "/ 1  # Fixed division by zero", user_code)
+        
+        # AI-powered optimization
+        prompt = f"Optimize this {lang} code for efficiency and security:\n```{lang}\n{user_code}\n```\nOptimized version:"
+        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        
+        with torch.no_grad():
+            outputs = model.generate(**inputs, max_length=1024)
+        
+        optimized_code = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        
+        # Extract just the code block if LLM added explanation
+        code_match = re.search(r'```(?:python)?\n(.*?)\n```', optimized_code, re.DOTALL)
+        if code_match:
+            optimized_code = code_match.group(1)
+        
+        return optimized_code if optimized_code else user_code
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=f"AI optimization failed: {str(e)}")
 
 # --- API Endpoints ---
 @app.post("/evaluate")
@@ -148,7 +167,11 @@ async def optimize_endpoint(request: CodeRequest):
 
 @app.get("/")
 def health_check():
-    return {"status": "Code Evaluation API is running!"}
+    return {
+        "status": "Code Evaluation API is running!",
+        "model_loaded": MODEL_NAME,
+        "cache_dir": str(CACHE_DIR)
+    }
 
 # For local testing
 if __name__ == "__main__":