Update api.py

api.py

@@ -1,31 +1,28 @@
-import time
 from flask import Flask, request, jsonify
 from app import generate_text
 
 app = Flask(__name__)
 
-# Warm up the model on startup
-print("Warming up model...")
-try:
-    generate_text("Warm up", max_length=10)
-    print("Model ready!")
-except Exception as e:
-    print(f"Warmup failed: {str(e)}")
-
-@app.route("/generate", methods=["POST"])
+@app.route('/generate', methods=['POST'])
 def generate():
     data = request.json
-    prompt = data.get("prompt", "")
-    max_length = data.get("max_length", 100)
+    prompt = data.get('prompt', '')
     
     if not prompt:
         return jsonify({"error": "Prompt is required"}), 400
     
     try:
-        result = generate_text(prompt, max_length)
-        return jsonify({"generated_text": result})
+        result = generate_text(
+            prompt,
+            max_new_tokens=data.get('max_new_tokens', 560),
+            max_context=data.get('max_context', 1080)
+        )
+        return jsonify({
+            "generated_text": result,
+            "tokens_used": len(prompt.split()) + len(result.split())
+        })
     except Exception as e:
         return jsonify({"error": str(e)}), 500
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=8080)
+if __name__ == '__main__':
+    app.run(host='0.0.0.0', port=8080)