Update main.py

main.py

@@ -1,11 +1,178 @@
 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
 
+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
+
 app = FastAPI()
 
+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()
+print('初始化完成')
+
 pipe_flan = pipeline("text2text-generation", model="google/flan-t5-small")
 
 @app.get("/infer_t5")