Spaces:
Running
Running
Update main.py
Browse files
main.py
CHANGED
|
@@ -6,7 +6,13 @@ from fastapi import FastAPI
|
|
| 6 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
| 7 |
|
| 8 |
model_name = "facebook/blenderbot-400M-distill"
|
|
|
|
|
|
|
|
|
|
|
|
|
| 9 |
# facebook/blenderbot-400M-distill
|
|
|
|
|
|
|
| 10 |
|
| 11 |
app = FastAPI()
|
| 12 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
|
|
@@ -14,7 +20,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
|
|
| 14 |
|
| 15 |
class req(BaseModel):
|
| 16 |
prompt: str
|
| 17 |
-
length: int
|
| 18 |
|
| 19 |
@app.get("/")
|
| 20 |
def read_root():
|
|
@@ -23,8 +28,6 @@ def read_root():
|
|
| 23 |
@app.post("/api")
|
| 24 |
def read_root(data: req):
|
| 25 |
print("Prompt:", data.prompt)
|
| 26 |
-
print("Length:", data.length)
|
| 27 |
-
print("Generating")
|
| 28 |
|
| 29 |
input_text = data.prompt
|
| 30 |
|
|
|
|
| 6 |
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
|
| 7 |
|
| 8 |
model_name = "facebook/blenderbot-400M-distill"
|
| 9 |
+
|
| 10 |
+
# https://huggingface.co/models?sort=trending&search=facebook%2Fblenderbo
|
| 11 |
+
# facebook/blenderbot-3B
|
| 12 |
+
# facebook/blenderbot-1B-distill
|
| 13 |
# facebook/blenderbot-400M-distill
|
| 14 |
+
# facebook/blenderbot-90M
|
| 15 |
+
# facebook/blenderbot_small-90M
|
| 16 |
|
| 17 |
app = FastAPI()
|
| 18 |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
|
|
|
|
| 20 |
|
| 21 |
class req(BaseModel):
|
| 22 |
prompt: str
|
|
|
|
| 23 |
|
| 24 |
@app.get("/")
|
| 25 |
def read_root():
|
|
|
|
| 28 |
@app.post("/api")
|
| 29 |
def read_root(data: req):
|
| 30 |
print("Prompt:", data.prompt)
|
|
|
|
|
|
|
| 31 |
|
| 32 |
input_text = data.prompt
|
| 33 |
|