Update app.py

app.py

@@ -1,45 +1,62 @@
-import flask # Gradio nahi, Flask hi rakhte hain
+import flask
 from flask import request, jsonify
-# from transformers import pipeline # अब इसकी ज़रूरत नहीं
-# import torch # अब इसकी ज़रूरत नहीं
+from transformers import pipeline
+import torch
+import warnings # warning suppress करने के लिए
 
-from ctransformers import AutoModelForCausalLM # ctransformers से मॉडल लोड करेंगे
+# warnings को suppress करें, वर्ना CPU पर warnings आ सकती हैं
+warnings.filterwarnings("ignore")
 
 app = flask.Flask(__name__)
 
 # ===========================
-# LOAD MODEL
+# LOAD MODEL (StableLM-3B-Chat)
 # ===========================
-model_id = "Qwen/Qwen1.5-1.8B-Chat-GGUF"
+model_id = "stabilityai/StableLM-3B-4E1T-Chat" 
 print("🔄 Loading model...")
 
+# CPU/GPU device set
+# हम CPU पर लोड करते समय 'torch.bfloat16' का उपयोग करके मेमोरी को कम करने की कोशिश करेंगे।
+device = 0 if torch.cuda.is_available() else -1
+dtype = torch.float32 if device == -1 else torch.bfloat16 # CPU के लिए float32
+
 try:
-    # ctransformers का उपयोग करके GGUF मॉडल को CPU पर लोड करें
-    ai = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        model_file="qwen1_5-1_8b-chat-q5_k_m.gguf", # GGUF फ़ाइल का नाम
-        model_type="qwen",
-        gpu_layers=0 # CPU पर चलाने के लिए
+    ai = pipeline(
+        "text-generation", 
+        model=model_id, 
+        max_new_tokens=200, 
+        device=device,
+        torch_dtype=dtype, # CPU/Memory optimization
+        trust_remote_code=True # StableLM के लिए आवश्यक
     )
     print("✅ Model loaded!")
 except Exception as e:
     print(f"❌ Error loading model: {e}")
-    # Fallback/Exit strategy here if loading fails
+    ai = None # If load fails, prevent later API errors
 
 # ===========================
 # CHAT API
 # ===========================
 @app.route('/chat', methods=['POST'])
 def chat():
+    if ai is None:
+        return jsonify({"error": "Model initialization failed."}), 500
+        
     try:
         data = request.get_json()
         msg = data.get("message", "")
         if not msg:
             return jsonify({"error": "No message sent"}), 400
 
-        # ctransformers से response generate करें
-        output = ai(msg, max_new_tokens=200, temperature=0.7)
-        return jsonify({"reply": output})
+        # StableLM Instruction Format:
+        prompt = f"<|user|>\n{msg}<|end|>\n<|assistant|>"
+        
+        output = ai(prompt)[0]["generated_text"]
+        
+        # Output को clean करें ताकि सिर्फ assistant का जवाब मिले
+        reply = output.split("<|assistant|>")[-1].strip()
+        
+        return jsonify({"reply": reply})
     except Exception as e:
         return jsonify({"error": str(e)}), 500