from flask import Flask, request, render_template
import transformers
import parrot
import torch
#parrot = Parrot(model_tag="prithivida/parrot_paraphraser_on_T5")

app = Flask(__name__)

@app.route("/", methods=["GET", "POST"])
def index():
    if request.method == "POST":
        text = request.form["text"]
        library = request.form["library"]

        if library == "transformers":
            model = transformers.AutoModelWithLMHead.from_pretrained("gpt2")
            paraphrased_text = transformers.generate(model, text)
        elif library == "parrot":
            paraphrased_text = parrot.augment(input_phrase=text, use_gpu=False)
        else:
            return "Error: Invalid library selected"

        return render_template("index.html", original_text=text, paraphrased_text=paraphrased_text)

    return render_template("index.html")

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