class Chat(chat_pb2_grpc.ChatServicer):
# WEB_ASK_06DEVBROS/ai/chatbot/checkpoint에 저장된 pth 파일(pytorch weight 파일)을 불러옴
root_path = str(
pathlib.Path(__file__).parent.absolute()) + '/../ai/chatbot'
checkpoint_path = f"{root_path}/checkpoint"
save_ckpt_path = f"{checkpoint_path}/kogpt2-wellness-auto-regressive.pth"
ctx = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(ctx)
# 저장한 Checkpoint 불러오기
checkpoint = torch.load(save_ckpt_path, map_location=device)
model = DialogKoGPT2()
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
tokenizer = get_kogpt2_tokenizer()
def ChatBot(self, request, context):
reqChat = request.clientChat #들어온 채팅 데이터
tokenized_indexs = self.tokenizer.encode(reqChat)
input_ids = torch.tensor([
self.tokenizer.bos_token_id,
] + tokenized_indexs + [self.tokenizer.eos_token_id]).unsqueeze(0)
output = self.model.generate(input_ids=input_ids)
resChat = self.tokenizer.decode(
output[0].tolist()[len(tokenized_indexs) + 1:],
skip_special_tokens=True)
return chat_pb2.ChatReply(serverChat=resChat)
def __init__(self, root_path='../ai/chatbot'):
checkpoint_path = f"{root_path}/checkpoint"
self.model_path = f"{checkpoint_path}/kogpt2-wellness-auto-regressive.pth"
checkpoint = torch.load(self.model_path, map_location=device)
self.model = DialogKoGPT2()
self.model.load_state_dict(checkpoint['model_state_dict'])
self.model.eval()
self.tokenizer = get_kogpt2_tokenizer()
data_path = "../data/ChatbotData.txt"
checkpoint_path = "../checkpoint"
save_ckpt_path = f"{checkpoint_path}/kogpt2-wellnesee-auto-regressive.pth"
n_epoch = 5 # Num of Epoch
batch_size = 2 # 배치 사이즈
device = "cuda" if torch.cuda.is_available() else "cpu"
save_step = 100 # 학습 저장 주기
learning_rate = 5e-5 # Learning Rate
dataset = WellnessAutoRegressiveDataset()
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True)
model = DialogKoGPT2()
loss_fct = torch.nn.CrossEntropyLoss(ignore_index=3)
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
losses = []
for epoch in range(n_epoch):
count = 0
with tqdm(total=len(train_loader), desc=f"Train({epoch})") as pbar:
for i, data in enumerate(train_loader):
optimizer.zero_grad()
data = torch.stack(
data) # list of Tensor로 구성되어 있기 때문에 list를 stack을 통해 변환해준다.
data = data.transpose(1, 0)
outputs = model(data, labels=data)