Update main.py

main.py

@@ -1,8 +1,15 @@
+import os
 from flask import Flask, request, jsonify
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
 app = Flask("Response API")
+size = "small"
+# microsoft/DialoGPT-small
+# microsoft/DialoGPT-medium
+# microsoft/DialoGPT-large
+
+os.environ["TRANSFORMERS_CACHE"] = r"./cache"
 
 # Load the Hugging Face GPT-2 model and tokenizer
 model = GPT2LMHeadModel.from_pretrained("microsoft/DialoGPT-medium")
@@ -28,4 +35,4 @@ def receive_data():
   print("Answered with:", answer_data)
   return jsonify(answer_data)
 
-app.run(debug=False, port=7860)
+app.run(debug=False, port=7860)