Update app.py

app.py

@@ -6,23 +6,20 @@ import torch
 app = flask.Flask(__name__)
 
 # ---------------------------
-# SMALL LLM MODEL (1–2 GB)
+# SUPER FAST SMALL MODEL
 # ---------------------------
-# Best small model: SmolLM-1.7B-Chat
-model_id = "HuggingFaceTB/SmolLM2-1.7B-Instruct"
+model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
 
 print(f"🔄 Loading {model_id} model...")
 
-# Load tokenizer
 tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-# Load model (auto dtype to avoid errors)
+# Best dtype for CPU speed = bfloat16
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
 )
 
-# Device setup
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
@@ -31,7 +28,7 @@ print(f"✅ {model_id} loaded successfully!")
 # ---------------------------
 # Chat Endpoint
 # ---------------------------
-@app.route('/chat', methods=['POST'])
+@app.route("/chat", methods=["POST"])
 def chat():
     try:
         data = request.get_json()
@@ -40,25 +37,23 @@ def chat():
         if not msg:
             return jsonify({"error": "No message sent"}), 400
 
-        # SmolLM uses normal text prompt (no ChatML)
-        prompt = f"<|user|>\n{msg}\n<|assistant|>\n"
+        # SmolLM2 format: no chat special tokens needed
+        prompt = f"User: {msg}\nAssistant:"
 
         inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
         output = model.generate(
             **inputs,
-            max_new_tokens=256,
-            do_sample=True,
-            temperature=0.6,
-            top_p=0.8,
+            max_new_tokens=128,   # fast
+            do_sample=False,      # FASTEST
             pad_token_id=tokenizer.eos_token_id,
         )
 
         reply = tokenizer.decode(output[0], skip_special_tokens=True)
 
-        # Extract only assistant part
-        if "<|assistant|>" in reply:
-            reply = reply.split("<|assistant|>")[-1].strip()
+        # Remove prompt text from output
+        if "Assistant:" in reply:
+            reply = reply.split("Assistant:")[-1].strip()
 
         return jsonify({"reply": reply})
 
@@ -67,4 +62,4 @@ def chat():
 
 
 if __name__ == "__main__":
-    app.run(host='0.0.0.0', port=7860)
+    app.run(host="0.0.0.0", port=7860)