Update main.py

main.py

@@ -1,36 +1,40 @@
 import os
-from flask import Flask, request, jsonify
+from flask import Flask, request, jsonify, render_template
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
 app = Flask("Response API")
-size = "small"
+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("microsoft/DialoGPT-medium")
-tokenizer = GPT2Tokenizer.from_pretrained("microsoft/DialoGPT-medium")
+model = GPT2LMHeadModel.from_pretrained(name)
+tokenizer = GPT2Tokenizer.from_pretrained(name)
 
 @app.route("/", methods=["POST"])
 def receive_data():
   data = request.get_json()
   
-  print("Prompt:", data['prompt'])
-  print("Length:", data['length'])
+  print("Prompt:", data["prompt"])
+  print("Length:", data["length"])
   
-  input_text = data['prompt']
+  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)
+  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 receive_data():
+  return render_template("index.html")
+
 app.run(debug=False, port=7860)