모델 학습 및 테스트

강태원·2023년 12월 6일
1

학습

모델 파라미터 설정

# config.py
from transformers import AutoModelForSeq2SeqLM,AutoTokenizer

# 모델 학습 후에 가중치 저장할 폴더
model_path='./saved_model'
# 모델 초기 가중치 로드할 곳
model_name = "gogamza/kobart-base-v2"
# 데이터셋 파일
data_root='./dataset'

# 내가 학습한 모델 가중치의 유무에 따라 분기 처리
if os.path.exists(f'{model_path}/pytorch_model.bin'):
  print("Use Customized Model")
  model = AutoModelForSeq2SeqLM.from_pretrained(model_path)
else:
  print("Use Pretrained Model")
  model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)

# Training Arguments
args={
  num_train_epochs=3
  per_device_train_batch_size=32
  per_device_eval_batch_size=32
  
  # 저장 폴더에 가중치 덮어쓰기 하는 옵션
  overwrite_output_dir=True

  # 몇 번의 학습 스텝이 지난 후에 평가를 수행할지 결정
  # 하나의 배치를 처리했을 때가 하나의 스텝이라 볼 수 있음
  eval_steps=3000 

  # 몇 번의 학습 스텝이 지난 후에 체크포인트를 진행할 지 결정
  # 중간에 학습이 끊기더라도 다시 진행하기 용이함
  save_steps=3000

  # LR(Learning Rate) 스케쥴러에서 사용할 스텝 수
  # Warmup을 통해 LR을 조절해 파라미터 업데이트의 폭을 조절함
  warmup_steps=300

  # 평가 수행 주기를 Epoch와 Steps중 Steps로 선택
  evaluation_strategy="steps"
  
  # 메모리 관리를 위해 예측값에 대한 손실만 계산하게 하는 옵션
  prediction_loss_only=True
  
  # 학습 중 저장되는 체크포인트의 최대 개수
  save_total_limit=3
}

모델 학습 코드

# train.py
from transformers import Seq2SeqTrainingArguments,Seq2SeqTrainer,\
                         DataCollatorForSeq2Seq
import dataset
import config

train_dataset,test_dataset = dataset.make_dataset(
    config.data_root, config.tokenizer
)

# 각 배치를 모델 학습에 알맞은 형태로 바꿔주는 역할
data_collator = DataCollatorForSeq2Seq(
    tokenizer=config.tokenizer, model=config.model
)

training_args = Seq2SeqTrainingArguments(
    **config.args
    output_dir=config.model_path,
    )

trainer = Seq2SeqTrainer(
    model=config.model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
)

# 모델 학습 진행
try:
  trainer.train()
except Exception as e:
  print(f"Failed to train model caused by {e}")
  
try:
  trainer.save_model(config.model_path)
  print("Model saved successfully.")
except Exception as e:
  print(f"Failed to save model caused by {e}")

예전에 작성 했던 코드들을 다시 되돌아보며 velog에 작성하고 있는데,
왜 이랬지? 싶은 부분들은 수정해서 업로드하고 있지만 기능적으로 빠진 부분들은 당장 추가하기에 어렵다.

또한 모델 학습 중 Evaluation 되는 수치들에 대한 정보가 없어 이를 참고용으로 올리지 못하는 것도 아쉽다.

다음에 AI로 학습할 일이 있다면 학습 결과는 꼭 따로 보관해야겠다..

(현재 폴더 구조)

AI Folder
ㄴ config.py
ㄴ train.py
ㄴ saved_model/
   ㄴ (saved_model_weights_files)
ㄴ dataset/
   ㄴ data.tsv

위 코드까지 문제없이 돌아갔다면, saved_model 폴더에 가중치가 잘 저장되었을 것이다!

테스트

표준어로 바뀐 사투리를 영어로도 바꾸기 위해 라이브러리를 하나 설치해주자!

pip install deep-translater
# test.py
from transformers import pipeline
from deep_translator import GoogleTranslator

nlg_pipeline=pipeline('text2text-generation',model=config.model_path,tokenizer=config.model_name)
def generate_text(pipe, text, num_return_sequences, max_length):
  target_style_name = "표준어"
  text = f"{target_style_name} 말투로 변환:{text}"
  out = pipe(text, num_return_sequences=num_return_sequences, max_length=max_length)
  return [x['generated_text'] for x in out]

num_return_sequences의 값에 따라서 반환되는 텍스트의 개수가 바뀐다
만약 3으로 지정했다면 길이가 3인 리스트에 담겨서 값이 반환될 것임!

def get_koen_text(target_text_ko):
  return GoogleTranslator(source='ko',target='en').translate(target_text_ko)
print("Write 'q' to exit")
while True:
  src_text=input("Dialect to translate : ")
  if src_text == 'q':
    break
  target_text_ko=generate_text(nlg_pipeline,src_text,num_return_sequences=1,max_length=64)[0]
  target_text_en = get_koen_text(target_text_ko)
  print(f"Translated Standard : {target_text_ko}")
  print(f"Translated Dialect : {target_text_en}")

여기까지 잘 따라와졌다면, 여러분은 경상도 사투리를 표준어로 번역해서 영어로 보여줄 수 있는 프로그램을 작성한 것이다! 아직 시도해보지는 않았지만 AI HUB에 있는 다른 방언 데이터셋을 이용해서 비슷한 방법으로 학습한다면 다른 사투리도 충분히 표준어로 바꿀 수 있다고 생각한다.

profile
가치를 창출하는 개발자! 가 목표입니다

0개의 댓글