Update main.py

main.py

@@ -19,6 +19,7 @@ import uvicorn
 from pydantic import BaseModel
 from transformers import pipeline
 
+extra_args = {}
 def set_args():
     """
     Sets up the arguments.
@@ -161,21 +162,18 @@ 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]")
+    # args = set_args()
+    acuda = torch.cuda.is_available() and not args.no_cuda
+    device = 'cuda' if acuda else 'cpu'
+    os.environ["CUDA_VISIBLE_DEVICES"] = args.device
+    tokenizer = BertTokenizerFast(vocab_file='/app/bert-base-zh/vocab.txt', sep_token="[SEP]", pad_token="[PAD]", cls_token="[CLS]")
     # # tokenizer = BertTokenizer(vocab_file=args.voca_path)
-    # model = Word_BERT()
+    model = Word_BERT()
     # # model = model.load_state_dict(torch.load(args.model_path))
     # model = model.to(device)
-    # model.eval()
-    # return model
-    return None
+    model.eval()
+    return model
+    # return None
 
 model11 = model_init()
 print(os.getcwd())