Update app.py

app.py

@@ -1,25 +1,36 @@
 from flask import Flask, request, jsonify
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM
 import torch
 
 app = Flask(__name__)
 
+# Load TinyLlama model
 model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, device_map="auto")
 
 @app.route("/api/ask", methods=["POST"])
 def ask():
     data = request.get_json()
     prompt = data.get("prompt", "")
-    inputs = tokenizer(prompt, return_tensors="pt")
-    outputs = model.generate(**inputs, max_new_tokens=100)
-    reply = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    return jsonify({"reply": reply})
-
-@app.route("/")
-def home():
-    return "✅ TinyLlama API is running!"
+
+    # Make it respond like a chatbot
+    chat_prompt = f"### Instruction:\nYou are a helpful, friendly chatbot named Acla. Reply conversationally.\n\n### Input:\n{prompt}\n\n### Response:"
+
+    inputs = tokenizer(chat_prompt, return_tensors="pt").to(model.device)
+    outputs = model.generate(
+        **inputs,
+        max_new_tokens=150,
+        temperature=0.7,
+        top_p=0.9,
+        do_sample=True
+    )
+
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Strip off the system prompt from output
+    response = response.split("### Response:")[-1].strip()
+
+    return jsonify({"reply": response})
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)