Update app.py

app.py

@@ -1,23 +1,34 @@
+import os
+from pathlib import Path
 from flask import Flask, request, jsonify
 from flask_cors import CORS
 from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
+# 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)
 
 # Model configuration
 MODEL_NAME = "deepseek-ai/deepseek-r1-6b-chat"
-MAX_NEW_TOKENS = 512
-DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+MAX_NEW_TOKENS = 256  # Reduced for free tier limits
+DEVICE = "cpu"  # Force CPU for Hugging Face Spaces
 
-# Initialize model and tokenizer
+# Initialize model
 try:
-    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+    tokenizer = AutoTokenizer.from_pretrained(
+        MODEL_NAME,
+        cache_dir=str(cache_dir)
+    
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
+        cache_dir=str(cache_dir),
         device_map="auto",
-        torch_dtype=torch.bfloat16 if DEVICE == "cuda" else torch.float32
+        torch_dtype=torch.float32,
+        low_cpu_mem_usage=True
     )
     print("Model loaded successfully!")
 except Exception as e:
@@ -25,39 +36,64 @@ except Exception as e:
     model = None
 
 def generate_response(prompt):
-    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=MAX_NEW_TOKENS,
-        do_sample=True,
-        temperature=0.7,
-        top_p=0.9,
-        pad_token_id=tokenizer.eos_token_id
-    )
-    return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    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
+        )
+        return tokenizer.decode(outputs[0], skip_special_tokens=True)
+    except Exception as e:
+        return f"Error generating response: {str(e)}"
 
 @app.route('/chat', methods=['POST'])
 def chat():
     if not model:
         return jsonify({"error": "Model not loaded"}), 500
     
-    data = request.json
-    prompt = data.get("prompt", "")
+    data = request.get_json()
+    if not data or 'prompt' not in data:
+        return jsonify({"error": "No prompt provided"}), 400
     
+    prompt = data['prompt'].strip()
     if not prompt:
-        return jsonify({"error": "No prompt provided"}), 400
+        return jsonify({"error": "Empty prompt"}), 400
     
     try:
         response = generate_response(prompt)
+        # Clean up extra text after the final answer
+        response = response.split("</s>")[0].strip()
         return jsonify({"response": response})
-    
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
 @app.route('/health', methods=['GET'])
 def health_check():
-    status = "ready" if model else "unavailable"
-    return jsonify({"status": status})
+    status = {
+        "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"
+    }
+    return jsonify(status)
+
+@app.route('/')
+def home():
+    return jsonify({
+        "service": "DeepSeek Chat API",
+        "endpoints": {
+            "POST /chat": "Process chat prompts",
+            "GET /health": "Service health check"
+        },
+        "config": {
+            "max_tokens": MAX_NEW_TOKENS,
+            "model": MODEL_NAME
+        }
+    })
 
 if __name__ == '__main__':
     app.run(host='0.0.0.0', port=5000)