fixing app.py

app.py

@@ -1,34 +1,35 @@
 import os
 import torch
 from fastapi import FastAPI, HTTPException
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
 
-# Set cache directory
+# Set a writable cache directory
 os.environ["HF_HOME"] = "/tmp/huggingface"
 os.environ["TRANSFORMERS_CACHE"] = "/tmp/huggingface"
 
 # Model setup
-MODEL_NAME = "deepseek-ai/deepseek-llm-7b-base"
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MODEL_NAME = "google/gemma-2b"  # Smaller, CPU-friendly model
+DEVICE = "cpu"
 
-# Load 4-bit quantized model (for speed & efficiency)
-bnb_config = BitsAndBytesConfig(
-    load_in_4bit=True,  # Enable 4-bit inference
+# 4-bit Quantization for CPU
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
     bnb_4bit_compute_dtype=torch.float16,
-    bnb_4bit_use_double_quant=True,
+    bnb_4bit_use_double_quant=True
 )
 
+# Load model & tokenizer
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME,
-    quantization_config=bnb_config,
-    device_map="auto",
-    attn_implementation="flash_attention_2"  # Enables Flash Attention
+    MODEL_NAME, 
+    quantization_config=quantization_config,
+    device_map="cpu"
 )
 
-# Compile for even faster inference (PyTorch 2.0+)
-model = torch.compile(model)
+# Set generation config
+model.generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
+model.generation_config.pad_token_id = model.generation_config.eos_token_id
 
 # FastAPI app
 app = FastAPI()
@@ -36,27 +37,14 @@ app = FastAPI()
 # Request payload
 class TextGenerationRequest(BaseModel):
     prompt: str
-    max_tokens: int = 512  # Default to 512
+    max_tokens: int = Field(default=100, ge=1, le=512)  # Prevent too large token requests
 
 @app.post("/generate")
 async def generate_text(request: TextGenerationRequest):
     try:
-        inputs = tokenizer(request.prompt, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
-
-        with torch.no_grad():
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=request.max_tokens,
-                do_sample=True,
-                temperature=0.7,
-                top_k=50,
-                top_p=0.9,
-                repetition_penalty=1.05,
-                use_cache=True,
-            )
-
+        inputs = tokenizer(request.prompt, return_tensors="pt").to(DEVICE)
+        outputs = model.generate(**inputs, max_new_tokens=request.max_tokens, do_sample=True)
         result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"generated_text": result}
-
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))