def get_model_and_tokenizer(model_name, device):
save_ckpt_path = CHECK_POINT[model_name]
if model_name == "koelectra":
model_name_or_path = "monologg/koelectra-base-discriminator"
tokenizer = ElectraTokenizer.from_pretrained(model_name_or_path)
electra_config = ElectraConfig.from_pretrained(model_name_or_path)
model = koElectraForSequenceClassification.from_pretrained(
pretrained_model_name_or_path=model_name_or_path,
config=electra_config,
num_labels=359)
elif model_name == 'kobert':
tokenizer = get_tokenizer()
model = KoBERTforSequenceClassfication()
if os.path.isfile(save_ckpt_path):
checkpoint = torch.load(save_ckpt_path, map_location=device)
pre_epoch = checkpoint['epoch']
# pre_loss = checkpoint['loss']
model.load_state_dict(checkpoint['model_state_dict'])
print(f"load pretrain from: {save_ckpt_path}, epoch={pre_epoch}")
return model, tokenizer
def __init__(self):
self.root_path = '..'
self.checkpoint_path = f"{self.root_path}/checkpoint"
self.save_ckpt_path = f"{self.checkpoint_path}/kobert-wellnesee-text-classification.pth"
#답변과 카테고리 불러오기
self.category, self.answer = load_wellness_answer()
ctx = "cuda" if torch.cuda.is_available() else "cpu"
self.device = torch.device(ctx)
# 저장한 Checkpoint 불러오기
checkpoint = torch.load(self.save_ckpt_path, map_location=self.device)
self.model = KoBERTforSequenceClassfication()
self.model.load_state_dict(checkpoint['model_state_dict'])
self.model.eval()
self.tokenizer = get_tokenizer()
class DialogKoBERT:
def __init__(self):
self.root_path = '..'
self.checkpoint_path = f"{self.root_path}/checkpoint"
self.save_ckpt_path = f"{self.checkpoint_path}/kobert-wellnesee-text-classification.pth"
#답변과 카테고리 불러오기
self.category, self.answer = load_wellness_answer()
ctx = "cuda" if torch.cuda.is_available() else "cpu"
self.device = torch.device(ctx)
# 저장한 Checkpoint 불러오기
checkpoint = torch.load(self.save_ckpt_path, map_location=self.device)
self.model = KoBERTforSequenceClassfication()
self.model.load_state_dict(checkpoint['model_state_dict'])
self.model.eval()
self.tokenizer = get_tokenizer()
def predict(self, sentence):
data = kobert_input(self.tokenizer, sentence, self.device, 512)
output = self.model(**data)
logit = output
softmax_logit = nn.Softmax(logit).dim
softmax_logit = softmax_logit[0].squeeze()
max_index = torch.argmax(softmax_logit).item()
max_index_value = softmax_logit[torch.argmax(softmax_logit)].item()
answer_list = self.answer[self.category[str(max_index)]]
answer_len = len(answer_list) - 1
answer_index = random.randint(0, answer_len)
return answer_list[answer_index]
if __name__ == "__main__":
root_path = '..'
checkpoint_path = f"{root_path}/checkpoint"
save_ckpt_path = f"{checkpoint_path}/kobert-wellnesee-text-classification.pth"
#답변과 카테고리 불러오기
category, answer = load_wellness_answer()
ctx = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(ctx)
# 저장한 Checkpoint 불러오기
checkpoint = torch.load(save_ckpt_path, map_location=device)
model = KoBERTforSequenceClassfication()
model.load_state_dict(checkpoint['model_state_dict'])
model.eval()
tokenizer = get_tokenizer()
while 1:
sent = input('\nQuestion: ') # '요즘 기분이 우울한 느낌이에요'
data = kobert_input(tokenizer, sent, device, 512)
# print(data)
output = model(**data)
logit = output
softmax_logit = nn.Softmax(logit).dim
checkpoint_path = "../checkpoint"
save_ckpt_path = f"{checkpoint_path}/kobert-wellnesee-text-classification.pth"
n_epoch = 20 # Num of Epoch
batch_size = 4 # 배치 사이즈
device = "cuda" if torch.cuda.is_available() else "cpu"
save_step = 100 # 학습 저장 주기
learning_rate = 5e-5 # Learning Rate
# WellnessTextClassificationDataset 데이터 로더
dataset = WellnessTextClassificationDataset()
train_loader = torch.utils.data.DataLoader(dataset,
batch_size=batch_size,
shuffle=True)
model = KoBERTforSequenceClassfication()
model.to(device)
# Prepare optimizer and schedule (linear warmup and decay)
no_decay = ['bias', 'LayerNorm.weight']
optimizer_grouped_parameters = [{
'params': [
p for n, p in model.named_parameters()
if not any(nd in n for nd in no_decay)
],
'weight_decay':
0.01
}, {
'params': [
p for n, p in model.named_parameters()
if any(nd in n for nd in no_decay)