Update main.py

main.py

@@ -103,7 +103,7 @@ class Word_BERT(nn.Module):
         ly_transfer = self.ly_transfer(pool)
         return out,cancer,transfer,ly_transfer
 
-def getChat(text,model):
+def getChat(text,model,tokenizer):
     # while True:
     #     if True:
     # text = input("user:")
@@ -172,14 +172,14 @@ def model_init():
     # # model = model.load_state_dict(torch.load(args.model_path))
     # model = model.to(device)
     model.eval()
-    return model
+    return tokenizer,model
     # return None
 
-model_extra = model_init()
+tokenizer,model_extra = model_init()
 
 @app.get("/infer_t5")
 def t5(input):
-    r1,r2,r3,r4 = getChat(input,model_extra)
+    r1,r2,r3,r4 = getChat(input,model_extra,tokenizer)
     # output = pipe_flan(input)
     return {"output": r4}