Update main.py

main.py

@@ -10,6 +10,8 @@ import torch
 
 app = FastAPI()
 name = "microsoft/DialoGPT-small"
+customGen = False
+gpt2based = False
 
 # microsoft/DialoGPT-small
 # microsoft/DialoGPT-medium
@@ -37,7 +39,7 @@ def read_root(data: req):
   print("Prompt:", data.prompt)
   print("Length:", data.length)
   
-  if name == "microsoft/DialoGPT-small" or name == "microsoft/DialoGPT-medium" or name == "microsoft/DialoGPT-large":
+  if (name == "microsoft/DialoGPT-small" or name == "microsoft/DialoGPT-medium" or name == "microsoft/DialoGPT-large") and customGen == True:
     # tokenizer = AutoTokenizer.from_pretrained("microsoft/DialoGPT-small")
     # model = AutoModelForCausalLM.from_pretrained("microsoft/DialoGPT-small")
     
@@ -58,16 +60,31 @@ def read_root(data: req):
     
     return answer_data
   else:
-    input_text = data.prompt
+    if gpt2based == True:
+      input_text = data.prompt
     
-    # Tokenize the input text
-    input_ids = gpt2tokenizer.encode(input_text, return_tensors="pt")
+      # Tokenize the input text
+      input_ids = gpt2tokenizer.encode(input_text, return_tensors="pt")
+      
+      # Generate output using the model
+      output_ids = gpt2model.generate(input_ids, max_length=data.length, num_beams=5, no_repeat_ngram_size=2)
+      generated_text = gpt2tokenizer.decode(output_ids[0], skip_special_tokens=True)
+      
+      answer_data = { "answer": generated_text }
+      print("Answer:", generated_text)
+      
+      return answer_data
+    else:
+      input_text = data.prompt
     
-    # 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 = gpt2tokenizer.decode(output_ids[0], skip_special_tokens=True)
-    
-    answer_data = { "answer": generated_text }
-    print("Answer:", generated_text)
-    
-    return answer_data
+      # 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("Answer:", generated_text)
+      
+      return answer_data