Spaces:
Running
Running
Update main.py
Browse files
main.py
CHANGED
@@ -9,7 +9,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
|
|
9 |
import torch
|
10 |
|
11 |
app = FastAPI()
|
12 |
-
name = "
|
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("
|
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
|