fixing app.py

app.py

@@ -2,27 +2,33 @@ import os
 import torch
 from fastapi import FastAPI, HTTPException
 from pydantic import BaseModel
-from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig
+from transformers import AutoTokenizer, AutoModelForCausalLM, GenerationConfig, BitsAndBytesConfig
 
-# Set a writable cache directory
+# Set 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"
-DTYPE = torch.float16 if DEVICE == "cuda" else torch.bfloat16
 
-# Load model and tokenizer
+# Load 4-bit quantized model (for speed & efficiency)
+bnb_config = BitsAndBytesConfig(
+    load_in_4bit=True,  # Enable 4-bit inference
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=True,
+)
+
 tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
 model = AutoModelForCausalLM.from_pretrained(
-    MODEL_NAME, torch_dtype=DTYPE, device_map="auto"
+    MODEL_NAME,
+    quantization_config=bnb_config,
+    device_map="auto",
+    attn_implementation="flash_attention_2"  # Enables Flash Attention
 )
 
-# Set up generation config
-generation_config = GenerationConfig.from_pretrained(MODEL_NAME)
-generation_config.pad_token_id = generation_config.eos_token_id
-generation_config.use_cache = True  # Speed up decoding
+# Compile for even faster inference (PyTorch 2.0+)
+model = torch.compile(model)
 
 # FastAPI app
 app = FastAPI()
@@ -30,28 +36,26 @@ app = FastAPI()
 # Request payload
 class TextGenerationRequest(BaseModel):
     prompt: str
-    max_tokens: int = 512  # Default to 512 for better performance
+    max_tokens: int = 512  # Default to 512
 
 @app.post("/generate")
 async def generate_text(request: TextGenerationRequest):
     try:
-        # Tokenize input and move tensors to the correct device
         inputs = tokenizer(request.prompt, return_tensors="pt", padding=True, truncation=True).to(DEVICE)
 
-        # Use no_grad() for faster inference
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=request.max_tokens,
-                do_sample=True,  # Enables sampling (use False for deterministic results)
-                temperature=0.7,  # Adjust for creativity (lower = more conservative)
-                top_k=50,  # Consider top 50 token choices
-                top_p=0.9,  # Nucleus sampling (reduces unlikely words)
-                repetition_penalty=1.1,  # Prevents looping responses
+                do_sample=True,
+                temperature=0.7,
+                top_k=50,
+                top_p=0.9,
+                repetition_penalty=1.05,
+                use_cache=True,
             )
 
-        # Decode generated tokens
-        result = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
+        result = tokenizer.decode(outputs[0], skip_special_tokens=True)
         return {"generated_text": result}
 
     except Exception as e: