Update app.py

app.py

@@ -16,7 +16,8 @@ app = Flask(__name__)
 CORS(app)  # Allow cross-origin requests
 
 # Model configuration
-MODEL_NAME = "deepseek-ai/deepseek-r1-6b-chat"
+# Use DeepSeek R1 Distill Qwen 7B model
+MODEL_NAME = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
 MAX_NEW_TOKENS = 256
 DEVICE = "cpu" if not torch.cuda.is_available() else "cuda"
 
@@ -32,24 +33,44 @@ def load_model():
         return True
     
     try:
-        from transformers import AutoTokenizer, AutoModelForCausalLM
+        from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig
         print(f"Loading model {MODEL_NAME}...")
         print(f"Using device: {DEVICE}")
         print(f"Cache directory: {cache_dir}")
         
+        # Use 4-bit quantization for memory efficiency if on CUDA
+        if DEVICE == "cuda":
+            quantization_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type="nf4",
+                bnb_4bit_use_double_quant=True
+            )
+        else:
+            quantization_config = None
+        
         # Load tokenizer
         tokenizer = AutoTokenizer.from_pretrained(
             MODEL_NAME,
-            cache_dir=str(cache_dir)
+            cache_dir=str(cache_dir),
+            trust_remote_code=True
         )
         
-        # Load model with low memory settings
+        # Configure token if HF_TOKEN is set
+        hf_token = os.environ.get("HF_TOKEN")
+        token_kwargs = {"token": hf_token} if hf_token else {}
+        
+        # Load model with appropriate settings for the device
         model = AutoModelForCausalLM.from_pretrained(
             MODEL_NAME,
             cache_dir=str(cache_dir),
             device_map="auto" if DEVICE == "cuda" else None,
             torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
-            low_cpu_mem_usage=True)
+            quantization_config=quantization_config,
+            low_cpu_mem_usage=True,
+            trust_remote_code=True,
+            **token_kwargs
+        )
         
         print("✅ Model loaded successfully!")
         return True
@@ -79,7 +100,15 @@ def stream_generator(prompt):
     
     # Prepare streaming generation
     try:
-        inputs = tokenizer(prompt, return_tensors="pt")
+        # Format prompt for the model
+        if "mistral" in MODEL_NAME.lower():
+            formatted_prompt = f"<s>[INST] {prompt} [/INST]"
+        elif "deepseek" in MODEL_NAME.lower():
+            formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        else:
+            formatted_prompt = prompt
+            
+        inputs = tokenizer(formatted_prompt, return_tensors="pt")
         if DEVICE == "cuda":
             inputs = inputs.to("cuda")
         
@@ -163,7 +192,15 @@ def chat():
         return jsonify({"error": "Empty prompt"}), 400
     
     try:
-        inputs = tokenizer(prompt, return_tensors="pt")
+        # Format prompt for the model
+        if "mistral" in MODEL_NAME.lower():
+            formatted_prompt = f"<s>[INST] {prompt} [/INST]"
+        elif "deepseek" in MODEL_NAME.lower():
+            formatted_prompt = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        else:
+            formatted_prompt = prompt
+            
+        inputs = tokenizer(formatted_prompt, return_tensors="pt")
         if DEVICE == "cuda":
             inputs = inputs.to("cuda")