api.py

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"])
def generate():
    data = request.json
    prompt = data.get("prompt", "")
    max_length = data.get("max_length", 100)
    
    if not prompt:
        return jsonify({"error": "Prompt is required"}), 400
    
    try:
        result = generate_text(prompt, max_length)
        return jsonify({"generated_text": result})
    except Exception as e:
        return jsonify({"error": str(e)}), 500

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8080)