from flask import Flask, render_template, request
from transformers import pipeline

app = Flask(__name__)

generator = pipeline(
    "text-generation",
    model="distilgpt2",
    device=-1
)

@app.route("/", methods=["GET", "POST"])
def home():
    response = ""
    if request.method == "POST":
        prompt = request.form["prompt"]
        result = generator(prompt, max_new_tokens=100)
        response = result[0]["generated_text"]

    return render_template("index.html", response=response)

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