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 PAD = '[PAD]' pad_id = 0 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('/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): # 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().to(device) input_ids = input_ids.unsqueeze(0) mask_input = torch.ones_like(input_ids).long().to(device) # 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) 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) return size_list,cancer,transfer,ly_transfer import requests import uvicorn from pydantic import BaseModel from fastapi import FastAPI app = FastAPI() # import intel_extension_for_pytorch as ipex 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() # history = [] Allhistory = {} print('初始化完成') import time lastReplyTime = 0 @app.post("/") async def get_Chat(context): global lastReplyTime tempReplyTime = int(time.time() * 1000) # if tempReplyTime % 10 == 0 or item1.must == True or tempReplyTime - lastReplyTime < 30000: # if item1.must == True: # lastReplyTime = tempReplyTime result = getChat(context) return {"res": result} @app.get("/hello") def read_root(): return {"Hello": "World!"}