Update app.py

app.py

@@ -1,9 +1,11 @@
 import os
+import time
+import json
 import numpy as np
 from pathlib import Path
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, Response
 from flask_cors import CORS
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 import torch
 
 # Verify numpy version
@@ -23,37 +25,90 @@ DEVICE = "cpu"
 
 # Initialize model
 try:
-    tokenizer = AutoTokenizer.from_pretrained(  # Fixed this line
+    tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         cache_dir=str(cache_dir)
-    )  # Added closing parenthesis
     
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         cache_dir=str(cache_dir),
         device_map="auto",
         torch_dtype=torch.float32,
-        low_cpu_mem_usage=True
-    )
+        low_cpu_mem_usage=True)
     print("Model loaded successfully!")
 except Exception as e:
     print(f"Model loading failed: {str(e)}")
     model = None
 
-def generate_response(prompt):
+def stream_generator(prompt):
+    """Generator function for streaming response with thinking steps"""
+    # Thinking phases
+    thinking_steps = [
+        "🔍 Analyzing your question...",
+        "🧠 Accessing knowledge base...",
+        "💡 Formulating response...",
+        "📚 Verifying information..."
+    ]
+    
+    # Stream thinking steps
+    for step in thinking_steps:
+        yield json.dumps({"type": "thinking", "content": step}) + '\n'
+        time.sleep(1.5)  # Simulate processing time
+    
+    # Prepare streaming generation
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+    streamer = TextStreamer(tokenizer, skip_prompt=True)
+    
+    # Generate response chunks
     try:
-        inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-        outputs = model.generate(
+        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
-        )
-        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+            pad_token_id=tokenizer.eos_token_id)
+        
+        # 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'
+                
     except Exception as e:
-        return f"Error generating response: {str(e)}"
+        yield json.dumps({
+            "type": "error",
+            "content": f"Generation error: {str(e)}"
+        }) + '\n'
+    
+    yield json.dumps({"type": "complete"}) + '\n'
+
+@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()
+    
+    if not prompt:
+        return jsonify({"error": "Empty prompt"}), 400
+    
+    return Response(
+        stream_generator(prompt),
+        mimetype='text/event-stream',
+        headers={
+            'Cache-Control': 'no-cache',
+            'Connection': 'keep-alive'
+        }
+    )
 
 @app.route('/chat', methods=['POST'])
 def chat():
@@ -61,18 +116,25 @@ def chat():
         return jsonify({"error": "Model not loaded"}), 500
     
     data = request.get_json()
-    if not data or 'prompt' not in data:
-        return jsonify({"error": "No prompt provided"}), 400
+    prompt = data.get('prompt', '').strip()
     
-    prompt = data['prompt'].strip()
     if not prompt:
         return jsonify({"error": "Empty prompt"}), 400
     
     try:
-        response = generate_response(prompt)
-        # Clean up extra text after the final answer
+        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)
+        
+        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         response = response.split("</s>")[0].strip()
         return jsonify({"response": response})
+    
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
@@ -82,7 +144,9 @@ def health_check():
         "model_loaded": bool(model),
         "device": DEVICE,
         "cache_dir": str(cache_dir),
-        "memory_usage": f"{torch.cuda.memory_allocated()/1024**2:.2f}MB" if torch.cuda.is_available() else "CPU"
+        "max_tokens": MAX_NEW_TOKENS,
+        "memory_usage": f"{torch.cuda.memory_allocated()/1024**2:.2f}MB" 
+            if torch.cuda.is_available() else "CPU"
     }
     return jsonify(status)
 
@@ -91,12 +155,14 @@ def home():
     return jsonify({
         "service": "DeepSeek Chat API",
         "endpoints": {
-            "POST /chat": "Process chat prompts",
+            "POST /chat": "Single-response chat",
+            "POST /stream_chat": "Streaming chat with thinking steps",
             "GET /health": "Service health check"
         },
         "config": {
+            "model": MODEL_NAME,
             "max_tokens": MAX_NEW_TOKENS,
-            "model": MODEL_NAME
+            "cache_location": str(cache_dir)
         }
     })