Update app.py

app.py

@@ -4,30 +4,28 @@ import torch
 
 app = Flask(__name__)
 
-# Load OpenChat model and tokenizer
-MODEL_NAME = "openchat/openchat-3.5-0106"
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, torch_dtype=torch.float16, device_map="auto")
+MODEL_NAME = "TheBloke/Mistral-7B-Instruct-v0.1-GPTQ"
 
-chatbot = pipeline("text-generation", model=model, tokenizer=tokenizer)
+tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, use_fast=True)
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_NAME,
+    device_map="auto",
+    torch_dtype=torch.float16,
+    trust_remote_code=True,
+    revision="gptq-4bit-32g-actorder_True"  # Based on the quant
+)
+
+generator = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 @app.route("/chat", methods=["POST"])
 def chat():
-    data = request.get_json()
-    prompt = data.get("message", "")
-    if not prompt:
+    user_input = request.json.get("message", "")
+    if not user_input:
         return jsonify({"error": "Empty message"}), 400
 
-    system_prompt = "<|system|>\nYou are a helpful assistant for food ordering.\n<|user|>\n" + prompt + "\n<|assistant|>\n"
-    output = chatbot(system_prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]["generated_text"]
-    
-    # Extract response after <|assistant|>
-    if "<|assistant|>" in output:
-        reply = output.split("<|assistant|>")[-1].strip()
-    else:
-        reply = output
-
-    return jsonify({"response": reply})
+    prompt = f"[INST] You are a helpful assistant for food ordering.\n{user_input} [/INST]"
+    result = generator(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)[0]['generated_text']
+    return jsonify({"response": result})
 
 if __name__ == "__main__":
-    app.run(debug=True)
+    app.run(host="0.0.0.0", port=7860)