Update app.py

app.py

@@ -5,44 +5,65 @@ import numpy as np
 from pathlib import Path
 from flask import Flask, request, jsonify, Response
 from flask_cors import CORS
-from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import torch
-
-# Verify numpy version
-assert np.__version__.startswith('1.'), f"Invalid numpy version {np.__version__} - must be 1.x series"
+import gc  # For garbage collection
 
 # Create cache directory if not exists
 cache_dir = Path(os.getenv('TRANSFORMERS_CACHE', '/app/cache'))
 cache_dir.mkdir(parents=True, exist_ok=True)
 
 app = Flask(__name__)
-CORS(app)
+CORS(app)  # Allow cross-origin requests
 
 # Model configuration
 MODEL_NAME = "deepseek-ai/deepseek-r1-6b-chat"
 MAX_NEW_TOKENS = 256
-DEVICE = "cpu"
+DEVICE = "cpu" if not torch.cuda.is_available() else "cuda"
 
-# Initialize model
-try:
-    tokenizer = AutoTokenizer.from_pretrained(
-        MODEL_NAME,
-        cache_dir=str(cache_dir)
-    )
+# Initialize model variables
+tokenizer = None
+model = None
+
+def load_model():
+    """Load model on first request to save memory at startup"""
+    global tokenizer, model
     
-    model = AutoModelForCausalLM.from_pretrained(
-        MODEL_NAME,
-        cache_dir=str(cache_dir),
-        device_map="auto",
-        torch_dtype=torch.float32,
-        low_cpu_mem_usage=True)
-    print("Model loaded successfully!")
-except Exception as e:
-    print(f"Model loading failed: {str(e)}")
-    model = None
+    if tokenizer is not None and model is not None:
+        return True
+    
+    try:
+        from transformers import AutoTokenizer, AutoModelForCausalLM
+        print(f"Loading model {MODEL_NAME}...")
+        print(f"Using device: {DEVICE}")
+        print(f"Cache directory: {cache_dir}")
+        
+        # Load tokenizer
+        tokenizer = AutoTokenizer.from_pretrained(
+            MODEL_NAME,
+            cache_dir=str(cache_dir)
+        )
+        
+        # Load model with low memory settings
+        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)
+        
+        print("✅ Model loaded successfully!")
+        return True
+    except Exception as e:
+        print(f"❌ Model loading failed: {str(e)}")
+        return False
 
 def stream_generator(prompt):
     """Generator function for streaming response with thinking steps"""
+    # Ensure model is loaded
+    if not load_model():
+        yield json.dumps({"type": "error", "content": "Model not loaded"}) + '\n'
+        return
+    
     # Thinking phases
     thinking_steps = [
         "🔍 Analyzing your question...",
@@ -54,48 +75,65 @@ def stream_generator(prompt):
     # Stream thinking steps
     for step in thinking_steps:
         yield json.dumps({"type": "thinking", "content": step}) + '\n'
-        time.sleep(1.5)  # Simulate processing time
+        time.sleep(0.8)  # Reduced timing for faster response
     
     # Prepare streaming generation
-    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-    streamer = TextStreamer(tokenizer, skip_prompt=True)
-    
-    # Generate response chunks
     try:
-        generated_ids = model.generate(
-            **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
-            streamer=streamer,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id)
+        inputs = tokenizer(prompt, return_tensors="pt")
+        if DEVICE == "cuda":
+            inputs = inputs.to("cuda")
         
-        # Stream generated text
-        full_response = ""
-        for token_ids in generated_ids:
-            chunk = tokenizer.decode(token_ids, skip_special_tokens=True)
-            new_content = chunk[len(full_response):]
-            if new_content.strip():
-                full_response = chunk
-                yield json.dumps({
-                    "type": "answer",
-                    "content": new_content
-                }) + '\n'
-                
+        # Use custom streaming implementation
+        # Start generation
+        with torch.no_grad():
+            generated_ids = model.generate(
+                **inputs,
+                max_new_tokens=MAX_NEW_TOKENS,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id,
+                return_dict_in_generate=True,
+                output_scores=False)
+        
+        # Get output sequence
+        output_ids = generated_ids.sequences[0][len(inputs.input_ids[0]):]
+        
+        # Stream in chunks for smoother experience
+        full_output = ""
+        chunk_size = 3  # Number of tokens per chunk
+        for i in range(0, len(output_ids), chunk_size):
+            chunk_ids = output_ids[i:i+chunk_size]
+            chunk_text = tokenizer.decode(chunk_ids, skip_special_tokens=True)
+            full_output += chunk_text
+            
+            yield json.dumps({
+                "type": "answer",
+                "content": chunk_text
+            }) + '\n'
+            
+            # Small delay for smoother streaming
+            time.sleep(0.05)
+            
     except Exception as e:
+        import traceback
+        error_details = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_details)
         yield json.dumps({
             "type": "error",
             "content": f"Generation error: {str(e)}"
         }) + '\n'
     
+    # Signal completion
     yield json.dumps({"type": "complete"}) + '\n'
+    
+    # Clean up memory
+    if DEVICE == "cuda":
+        torch.cuda.empty_cache()
+    gc.collect()
 
 @app.route('/stream_chat', methods=['POST'])
 def stream_chat():
-    if not model:
-        return jsonify({"error": "Model not loaded"}), 500
-    
     data = request.get_json()
     prompt = data.get('prompt', '').strip()
     
@@ -107,14 +145,16 @@ def stream_chat():
         mimetype='text/event-stream',
         headers={
             'Cache-Control': 'no-cache',
+            'X-Accel-Buffering': 'no',  # Prevent Nginx buffering
             'Connection': 'keep-alive'
         }
     )
 
 @app.route('/chat', methods=['POST'])
 def chat():
-    if not model:
-        return jsonify({"error": "Model not loaded"}), 500
+    # Ensure model is loaded
+    if not load_model():
+        return jsonify({"error": "Model failed to load"}), 500
     
     data = request.get_json()
     prompt = data.get('prompt', '').strip()
@@ -123,26 +163,48 @@ def chat():
         return jsonify({"error": "Empty prompt"}), 400
     
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=MAX_NEW_TOKENS,
-            temperature=0.7,
-            top_p=0.9,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id)
+        inputs = tokenizer(prompt, return_tensors="pt")
+        if DEVICE == "cuda":
+            inputs = inputs.to("cuda")
+            
+        with torch.no_grad():
+            outputs = model.generate(
+                **inputs,
+                max_new_tokens=MAX_NEW_TOKENS,
+                temperature=0.7,
+                top_p=0.9,
+                do_sample=True,
+                pad_token_id=tokenizer.eos_token_id)
+        
+        response = tokenizer.decode(outputs[0][len(inputs.input_ids[0]):], skip_special_tokens=True)
+        
+        # Clean up memory
+        if DEVICE == "cuda":
+            torch.cuda.empty_cache()
+        gc.collect()
         
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        response = response.split("</s>")[0].strip()
         return jsonify({"response": response})
     
     except Exception as e:
+        import traceback
+        error_details = f"Error: {str(e)}\n{traceback.format_exc()}"
+        print(error_details)
         return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
+    model_loaded = tokenizer is not None and model is not None
+    
+    try:
+        # Check if we need to load the model
+        if not model_loaded and request.args.get('load') == 'true':
+            model_loaded = load_model()
+    except Exception as e:
+        print(f"Health check error: {str(e)}")
+    
     status = {
-        "model_loaded": bool(model),
+        "status": "ok" if model_loaded else "waiting",
+        "model_loaded": model_loaded,
         "device": DEVICE,
         "cache_dir": str(cache_dir),
         "max_tokens": MAX_NEW_TOKENS,
@@ -155,6 +217,7 @@ def health_check():
 def home():
     return jsonify({
         "service": "DeepSeek Chat API",
+        "status": "online",
         "endpoints": {
             "POST /chat": "Single-response chat",
             "POST /stream_chat": "Streaming chat with thinking steps",
@@ -168,4 +231,9 @@ def home():
     })
 
 if __name__ == '__main__':
-    app.run(host='0.0.0.0', port=5000)
+    # Load model at startup - only if explicitly requested
+    if os.getenv('PRELOAD_MODEL', 'false').lower() == 'true':
+        load_model()
+    
+    port = int(os.environ.get("PORT", 5000))
+    app.run(host='0.0.0.0', port=port)