flan

question_paper.py

@@ -4,7 +4,7 @@ from fastapi.middleware.cors import CORSMiddleware
 from fastapi.encoders import jsonable_encoder
 from fastapi.responses import JSONResponse
 
-pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
+pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small", max_new_tokens=1024)
 
 app = FastAPI()
 
@@ -31,7 +31,7 @@ def llama():
     # tresponse = TOKENIZER.decode(tokens[0], skip_special_tokens=False)
     # print(tresponse)
     input = "Generate 5 jokes"
-    output = pipe_flan(input, {"max_new_tokens": 1024,})
+    output = pipe_flan(input)
     response_message = {"message": output[0]["generated_text"]}
     json_response = jsonable_encoder(response_message)
     return JSONResponse(content=json_response)