import os
from flask import Flask, request, jsonify, render_template
from transformers import GPT2LMHeadModel, GPT2Tokenizer
import torch

app = Flask("Response API")
name = "microsoft/DialoGPT-medium"
# microsoft/DialoGPT-small
# microsoft/DialoGPT-medium
# microsoft/DialoGPT-large

# Load the Hugging Face GPT-2 model and tokenizer
model = GPT2LMHeadModel.from_pretrained(name)
tokenizer = GPT2Tokenizer.from_pretrained(name)

@app.route("/api", methods=["POST"])
def receive_data():
  data = request.get_json()
  
  print("Prompt:", data["prompt"])
  print("Length:", data["length"])
  
  input_text = data["prompt"]
  
  # Tokenize the input text
  input_ids = tokenizer.encode(input_text, return_tensors="pt")
  
  # Generate output using the model
  output_ids = model.generate(input_ids, max_length=data["length"], num_beams=5, no_repeat_ngram_size=2)
  generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
  
  answer_data = { "answer": generated_text }
  print("Answered with:", answer_data)
  return jsonify(answer_data)

@app.route("/", methods=["GET"])
def not_api():
  return render_template("index.html")

app.run(debug=False, port=7860)