Update app.py

app.py

@@ -1,41 +1,48 @@
-from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
+from transformers import AutoModel, AutoTokenizer
+from flask import Flask, request, jsonify
 
 app = Flask(__name__)
 
-# Load model with memory optimizations
-model_name = "tanusrich/Mental_Health_Chatbot"
+# Model name from Hugging Face
+MODEL_NAME = "tanusrich/Mental_Health_Chatbot"
 
+# Detect if GPU is available
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
 # Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-
-# Load model with reduced memory usage
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",  # Automatically assign layers to CPU/GPU
-    torch_dtype=torch.float16,  # Use less memory with half-precision
-    low_cpu_mem_usage=True,  # Optimize RAM usage
-    max_memory={device: "4GB"}  # Limit model memory usage
-).to(device)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
+
+# Load model with optimized settings
+try:
+    model = AutoModel.from_pretrained(
+        MODEL_NAME,
+        device_map="auto",  # Automatically selects best available device
+        low_cpu_mem_usage=True  # Optimized for lower memory consumption
+    ).to(device)
+except Exception as e:
+    print(f"Error loading model: {e}")
+    exit(1)
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    user_input = request.json.get("message", "")
+    data = request.json
+    user_input = data.get("message", "")
 
     if not user_input:
         return jsonify({"error": "Message is required"}), 400
 
+    # Tokenize input
     inputs = tokenizer(user_input, return_tensors="pt").to(device)
-    
-    with torch.no_grad():  # Disable gradient calculation to save memory
-        output = model.generate(**inputs, max_length=100)
 
-    response = tokenizer.decode(output[0], skip_special_tokens=True)
-    
-    return jsonify({"response": response})
+    # Generate response
+    with torch.no_grad():
+        outputs = model.generate(**inputs, max_length=150)
+
+    # Decode response
+    response_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    return jsonify({"response": response_text})
 
 if __name__ == "__main__":
-    app.run(debug=True, port=5000)
+    app.run(host="0.0.0.0", port=5000)