from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse import torch import os import json import random import numpy as np from torch import nn import argparse import logging from transformers import GPT2TokenizerFast, GPT2LMHeadModel, GPT2Config from transformers import BertTokenizerFast import torch.nn.functional as F from transformers import AutoTokenizer, AutoConfig, get_linear_schedule_with_warmup, AdamW, BertModel import requests import uvicorn from pydantic import BaseModel from transformers import pipeline extra_args = {} def set_args(): """ Sets up the arguments. """ parser = argparse.ArgumentParser() parser.add_argument('--device', default='0', type=str, required=False, help='生成设备') # parser.add_argument('--model_config', default='config/model_config_dialogue_small.json', type=str, required=False, # help='模型参数') parser.add_argument('--log_path', default='interact.log', type=str, required=False, help='interact日志存放位置') parser.add_argument('--model_path', default='./pathology_extra/result/12/model.pth', type=str, required=False, help='对话模型路径') parser.add_argument('--vocab_path', default='/app/bert-base-zh/vocab.txt', type=str, required=False, help='对话模型路径') parser.add_argument('--repetition_penalty', default=1.0, type=float, required=False, help="重复惩罚参数，若生成的对话重复性较高，可适当提高该参数") # parser.add_argument('--seed', type=int, default=None, help='设置种子用于生成随机数，以使得训练的结果是确定的') parser.add_argument('--max_len', type=int, default=25, help='每个utterance的最大长度,超过指定长度则进行截断') parser.add_argument('--max_history_len', type=int, default=3, help="dialogue history的最大长度") parser.add_argument('--no_cuda', action='store_true', help='不使用GPU进行预测') return parser.parse_args() def create_logger(args): """ 将日志输出到日志文件和控制台 """ logger = logging.getLogger(__name__) logger.setLevel(logging.INFO) formatter = logging.Formatter( '%(asctime)s - %(levelname)s - %(message)s') # 创建一个handler，用于写入日志文件 file_handler = logging.FileHandler( filename=args.log_path) file_handler.setFormatter(formatter) file_handler.setLevel(logging.INFO) logger.addHandler(file_handler) # 创建一个handler，用于将日志输出到控制台 console = logging.StreamHandler() console.setLevel(logging.DEBUG) console.setFormatter(formatter) logger.addHandler(console) return logger class Word_BERT(nn.Module): def __init__(self, seq_label=1,cancer_label=8,transfer_label=2,ly_transfer=2): super(Word_BERT, self).__init__() self.bert = BertModel.from_pretrained('/home/user/app/./bert-base-zh') # self.bert_config = self.bert.config self.out = nn.Sequential( # nn.Linear(768,256), # nn.ReLU(), nn.Dropout(0.1), nn.Linear(768, seq_label) ) self.cancer = nn.Sequential( nn.Dropout(0.1), nn.Linear(768, cancer_label) ) self.transfer = nn.Sequential( nn.Dropout(0.1), nn.Linear(768, transfer_label) ) self.ly_transfer = nn.Sequential( nn.Dropout(0.1), nn.Linear(768, ly_transfer) ) def forward(self, word_input, masks): # print(word_input.size()) output = self.bert(word_input, attention_mask=masks) sequence_output = output.last_hidden_state pool = output.pooler_output # print(sequence_output.size()) # print(pool.size()) out = self.out(sequence_output) cancer = self.cancer(pool) transfer = self.transfer(pool) ly_transfer = self.ly_transfer(pool) return out,cancer,transfer,ly_transfer def getChat(text,model,tokenizer): # while True: # if True: # text = input("user:") # text = "你好" # if args.save_samples_path: # samples_file.write("user:{}\n".format(text)) text = ['[CLS]']+[i for i in text]+['[SEP]'] # print(text) text_ids = tokenizer.convert_tokens_to_ids(text) # print(text_ids) input_ids = torch.tensor(text_ids).long() input_ids = input_ids.unsqueeze(0) mask_input = torch.ones_like(input_ids).long() # print(input_ids.size()) response = [] # 根据context，生成的response # 最多生成max_len个token with torch.no_grad(): out, cancer, transfer, ly_transfer = model(input_ids, mask_input) out = F.sigmoid(out).squeeze(2).cpu() out = out.numpy().tolist() cancer = cancer.argmax(dim=-1).cpu().numpy().tolist() transfer = transfer.argmax(dim=-1).cpu().numpy().tolist() ly_transfer = ly_transfer.argmax(dim=-1).cpu().numpy().tolist() # print(out) # print(cancer,transfer,ly_transfer) pred_thresold = [[1 if jj > 0.4 else 0 for jj in ii] for ii in out] size_list = [] start,end = 0,0 for i,j in enumerate(pred_thresold[0]): if j==1 and start==end: start = i elif j!=1 and start!=end: end = i size_list.append((start,end)) start = end # print(size_list) size_text = [] for k in size_list: size_text.append(text[k[0]:k[1]]) if len(size_text)==0: size_str = "无" else: size_str = ''.join(size_text[0]) if '×' in size_str: split_w = '×' else: split_w = '*' tt = size_str.split(split_w) f = 0 if tt[0][0].isdigit(): size_float = [float(i) for i in tt] for kk in size_float: if kk>=4: f = 1 else: size_str = "无" if size_str == "无": size_4 = "无" elif f==0: size_4 = "<4cm" else: size_4 = ">=4cm" # if len(size_list)==0: # size_str = "无" # else: # size_str = ''.join(size_list) cancer_dict = {'腺癌': 0, '肺良性疾病': 1, '鳞癌': 2, '无法判断组织分型': 3, '复合型': 4, '转移癌': 5, '小细胞癌': 6, '大细胞癌': 7} id_cancer = {j:i for i,j in cancer_dict.items()} transfer_id = {'无': 0, '转移': 1} id_transfer = {j:i for i,j in transfer_id.items()} lymph_transfer_id = {'无': 0, '淋巴转移': 1} id_lymph_transfer = {j: i for i, j in lymph_transfer_id.items()} # print(cancer) cancer = id_cancer[cancer[0]] transfer = id_transfer[transfer[0]] ly_transfer = id_lymph_transfer[ly_transfer[0]] # print(cancer,transfer,ly_transfer) output = "肿瘤大小:"+size_str+"\n肿瘤大小<>=4cm:"+size_4+"\n"+"病理组织分型:"+cancer+"\n"+"转移:"+transfer+"\n"+"淋巴转移:"+ly_transfer+"\n" return output,size_str,size_4,cancer,transfer,ly_transfer app = FastAPI() def model_init(): # args = set_args() acuda = torch.cuda.is_available() and not args.no_cuda device = 'cuda' if acuda else 'cpu' os.environ["CUDA_VISIBLE_DEVICES"] = device tokenizer = BertTokenizerFast(vocab_file='/home/user/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.load_state_dict(torch.load('/home/user/app/./model.pth',map_location=torch.device('cpu'))) # model = model.to(device) model.eval() return tokenizer,model # return None tokenizer,model_extra = model_init() @app.get("/infer_t5") def t5(input): output,size_str,size_4,cancer,transfer,ly_transfer = getChat(input,model_extra,tokenizer) # output = pipe_flan(input) # ,"size_str":size_str,"size_4":size_4,"cancer":cancer,"transfer":transfer,"ly_transfer":ly_transfer return {"output": output,"size_str":size_str,"size_4":size_4,"cancer":cancer,"transfer":transfer,"ly_transfer":ly_transfer} app.mount("/", StaticFiles(directory="static", html=True), name="static") @app.get("/") def index() -> FileResponse: return FileResponse(path="/app/static/index.html", media_type="text/html") @app.get("/postText") def postText(input): output,size_str,size_4,cancer,transfer,ly_transfer = getChat(input,model_extra,tokenizer) return {"output": output,"size_str":size_str,"size_4":size_4,"cancer":cancer,"transfer":transfer,"ly_transfer":ly_transfer}