Update main.py

main.py

@@ -6,7 +6,13 @@ from fastapi import FastAPI
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 
 model_name = "facebook/blenderbot-400M-distill"
+
+# https://huggingface.co/models?sort=trending&search=facebook%2Fblenderbo
+# facebook/blenderbot-3B
+# facebook/blenderbot-1B-distill
 # facebook/blenderbot-400M-distill
+# facebook/blenderbot-90M
+# facebook/blenderbot_small-90M
 
 app = FastAPI()
 model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
@@ -14,7 +20,6 @@ tokenizer = AutoTokenizer.from_pretrained(model_name)
 
 class req(BaseModel):
   prompt: str
-  length: int
 
 @app.get("/")
 def read_root():
@@ -23,8 +28,6 @@ def read_root():
 @app.post("/api")
 def read_root(data: req):
   print("Prompt:", data.prompt)
-  print("Length:", data.length)
-  print("Generating")
 
   input_text = data.prompt