BeveledCube commited on
Commit
c858ace
·
verified ·
1 Parent(s): 20aacee

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +2 -2
main.py CHANGED
@@ -9,7 +9,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
9
  import torch
10
 
11
  app = FastAPI()
12
- name = "microsoft/DialoGPT-small"
13
 
14
  # microsoft/DialoGPT-small
15
  # microsoft/DialoGPT-medium
@@ -48,6 +48,6 @@ def read_root(data: req):
48
  generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
49
 
50
  answer_data = { "answer": generated_text }
51
- print("Answered with:", answer_data)
52
 
53
  return answer_data
 
9
  import torch
10
 
11
  app = FastAPI()
12
+ name = "PygmalionAI/pygmalion-1.3b"
13
 
14
  # microsoft/DialoGPT-small
15
  # microsoft/DialoGPT-medium
 
48
  generated_text = tokenizer.decode(output_ids[0], skip_special_tokens=True)
49
 
50
  answer_data = { "answer": generated_text }
51
+ print("Answer:", generated_text)
52
 
53
  return answer_data