Update main.py

main.py

@@ -9,6 +9,14 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 
 app = FastAPI()
+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)
 
 class req(BaseModel):
   prompt: str
@@ -20,15 +28,6 @@ def read_root():
 
 @app.post("/api")
 def read_root(data: req):
-  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)
-
   print("Prompt:", data.prompt)
   print("Length:", data.length)