Update server.py

server.py

@@ -3,33 +3,28 @@ from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
-import subprocess
-import tempfile
 
 app = FastAPI()
 
-# ✅ Fix: Use `/tmp` as cache directory (writable in Hugging Face Spaces)
+# ✅ Fix: Use `/tmp` as the cache directory (Hugging Face Spaces allows writing here)
 os.environ["TRANSFORMERS_CACHE"] = "/tmp"
 os.environ["HF_HOME"] = "/tmp"
 
-# ✅ Ensure cache directory exists
+# ✅ Ensure the `/tmp` directory exists
 if not os.path.exists("/tmp"):
     os.makedirs("/tmp")
 
-# ✅ Load AI Model
+# ✅ Load DeepSeek-Coder-V2-Base Model with correct cache directory
 model_name = "deepseek-ai/DeepSeek-Coder-V2-Base"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(model_name, cache_dir="/tmp")
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto", cache_dir="/tmp")
 
 class CodeRequest(BaseModel):
     user_story: str
 
-class TestRequest(BaseModel):
-    code: str
-
 @app.post("/generate-code")
 def generate_code(request: CodeRequest):
-    """Generates AI-powered structured code based on user story"""
+    """Generates structured AI-powered code based on user story"""
     prompt = f"Generate structured code for: {request.user_story}"
 
     inputs = tokenizer(prompt, return_tensors="pt").to("cuda" if torch.cuda.is_available() else "cpu")
@@ -38,37 +33,6 @@ def generate_code(request: CodeRequest):
 
     return {"generated_code": generated_code}
 
-@app.post("/test-code")
-def test_code(request: TestRequest):
-    """Runs automated testing on AI-generated code"""
-    try:
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".py", dir="/tmp") as temp_file:
-            temp_file.write(request.code.encode())
-            temp_file.close()
-
-            result = subprocess.run(["pytest", temp_file.name], capture_output=True, text=True)
-            os.unlink(temp_file.name)
-
-            if result.returncode == 0:
-                return {"test_status": "All tests passed!"}
-            else:
-                return {"test_status": "Test failed!", "details": result.stderr}
-
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/execute-code")
-def execute_code():
-    """Executes AI-generated code and returns output"""
-    sample_code = "print('Hello from AI-generated code!')"
-    
-    try:
-        result = subprocess.run(["python3", "-c", sample_code], capture_output=True, text=True)
-        return {"status": "Execution successful!", "output": result.stdout}
-    
-    except Exception as e:
-        return {"status": "Execution failed!", "error": str(e)}
-
 if __name__ == "__main__":
     import uvicorn
     uvicorn.run(app, host="0.0.0.0", port=7860)