updated model.py

app.py

@@ -5,17 +5,16 @@ from model import generate_code
 
 app = FastAPI()
 
-# Define request body schema
 class GenerateRequest(BaseModel):
     prompt: str
-    max_tokens: int = 256  # Default value
+    max_tokens: int = 256
 
 @app.get("/")
 def home():
     return {"message": "Code Generation API is running!"}
 
 @app.post("/generate")
-def generate(request: GenerateRequest):  # Expect JSON
+def generate(request: GenerateRequest):
     if not request.prompt:
         raise HTTPException(status_code=400, detail="Prompt cannot be empty.")
 

model.py

@@ -6,24 +6,23 @@ MODEL_NAME = "bigcode/starcoderbase-1b"
 HF_TOKEN = os.getenv("HUGGINGFACE_TOKEN")
 
 # Force CPU mode
-device = "cpu"  # Change this from "cuda"
+device = "cpu"
 
 # Load tokenizer and model
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, token=HF_TOKEN)
 
 # Ensure the tokenizer has a pad token set
 if tokenizer.pad_token is None:
-    tokenizer.pad_token = tokenizer.eos_token  # Set pad_token to eos_token
+    tokenizer.pad_token = tokenizer.eos_token
 
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_NAME,
     token=HF_TOKEN,
-    torch_dtype=torch.float16,  # Keep memory low
-    device_map="auto",  # Still allows auto placement
+    torch_dtype=torch.float32,  # Change to float32 for CPU compatibility
     trust_remote_code=True
-).to(device)
+).to(device)  # Explicitly move to CPU
 
 def generate_code(prompt: str, max_tokens: int = 256):
-    inputs = tokenizer(prompt, return_tensors="pt", padding=True).to(device)  # Enable padding
-    output = model.generate(**inputs, max_new_tokens=max_tokens, pad_token_id=tokenizer.pad_token_id)  # Explicit pad_token_id
+    inputs = tokenizer(prompt, return_tensors="pt", padding=True, truncation=True).to(device)
+    output = model.generate(**inputs, max_new_tokens=max_tokens, pad_token_id=tokenizer.pad_token_id)
     return tokenizer.decode(output[0], skip_special_tokens=True)