extra / main.py
qianmuuq's picture
Update main.py
78b9ab2
raw
history blame
7.67 kB
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='./bert-base-zh\\vocab.txt', type=str, required=False,
help='对话模型路径')
parser.add_argument('--save_samples_path', default="sample/", 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('./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: str, userid: int):
# 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('初始化完成')
# if __name__ == '__main__':
# # getChat("测试一下", 0)
# # main()
# uvicorn.run(app='main:app', host="localhost",
# port=7860, reload=False)
# # testFunc()
class Items1(BaseModel):
context: str
userid: int
# must: bool
import time
lastReplyTime = 0
@app.post("/")
async def get_Chat(item1: Items1):
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(
item1.context, item1.userid)
return {"res": result}
@app.get("/hello")
def read_root():
return {"Hello": "World!"}