Update main.py

main.py

@@ -160,6 +160,24 @@ app = FastAPI()
 
 pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
 
+def model_init():
+    args = set_args()
+    logger = create_logger(args)
+    # 当用户使用GPU,并且GPU可用时
+    args.cuda = torch.cuda.is_available() and not args.no_cuda
+    device = 'cuda' if args.cuda else 'cpu'
+    logger.info('using device:{}'.format(device))
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
+    tokenizer = BertTokenizerFast(vocab_file=args.vocab_path, sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
+    # tokenizer = BertTokenizer(vocab_file=args.voca_path)
+    model = Word_BERT()
+    # model = model.load_state_dict(torch.load(args.model_path))
+    model = model.to(device)
+    model.eval()
+    return model
+
+model = model_init()
+
 @app.get("/infer_t5")
 def t5(input):
     output = pipe_flan(input)