Update main.py

main.py

@@ -12,7 +12,7 @@ from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI(title="Code Evaluation & Optimization API")
 
-# Required CORS configuration
+# CORS Configuration
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -21,20 +21,20 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-# --- Environment Setup ---
+# 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 ---
+# Load AI Model
 MODEL_NAME = "codellama/CodeLlama-7b-hf"
 
 try:
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
-        cache_dir=str(CACHE_DIR)
-        model = AutoModelForCausalLM.from_pretrained(
+        cache_dir=str(CACHE_DIR))
+    model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         device_map="auto",
         torch_dtype=torch.float16,
@@ -42,12 +42,12 @@ try:
 except Exception as e:
     raise RuntimeError(f"Failed to load model: {str(e)}")
 
-# --- Request Model ---
+# Request Model
 class CodeRequest(BaseModel):
     code: str
-    language: str = "python"  # Default to Python
+    language: str = "python"
 
-# --- Helper Functions ---
+# Helper Functions
 def evaluate_code(user_code: str, lang: str) -> dict:
     """Evaluate code for correctness, performance, and security"""
     start_time = time.time()
@@ -131,7 +131,7 @@ def optimize_code_ai(user_code: str, lang: str) -> str:
     except Exception as e:
         raise HTTPException(status_code=500, detail=f"AI optimization failed: {str(e)}")
 
-# --- API Endpoints ---
+# API Endpoints
 @app.post("/evaluate")
 async def evaluate_endpoint(request: CodeRequest):
     try:
@@ -159,7 +159,6 @@ def health_check():
         }
     }
 
-# Required for Hugging Face Spaces
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run("main:app", host="0.0.0.0", port=7860)