Update app.py

app.py

@@ -1,28 +1,25 @@
 from flask import Flask, request, jsonify
-from transformers import AutoTokenizer, AutoModelForCausalLM
+from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 
 app = Flask(__name__)
 
-# Load model and tokenizer
-MODEL_NAME = "openai-community/gpt2"  # or your Tiny LLaMA model
-tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-model = AutoModelForCausalLM.from_pretrained(MODEL_NAME, device_map="auto")
+model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 @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=50)
-    text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-
-    return jsonify({"reply": text})
+    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 "✅ Model API running! POST JSON to /api/ask with {'prompt': 'your text'}"
+    return "✅ TinyLlama API is running!"
 
 if __name__ == "__main__":
     app.run(host="0.0.0.0", port=7860)