model_name = "SungBeom/whisper-small-ko" # 대안 : "SungBeom/whisper-small-ko"
dataset_name = "maxseats/aihub-464-preprocessed-680GB-set-0"  # 불러올 데이터셋(허깅페이스 기준)

CACHE_DIR = '/mnt/a/maxseats/.finetuning_cache'  # 캐시 디렉토리 지정
is_test = False  # True: 소량의 샘플 데이터로 테스트, False: 실제 파인튜닝

token = "hf_" # 허깅페이스 토큰 입력

training_args = Seq2SeqTrainingArguments(
    output_dir=model_dir,  # 원하는 리포지토리 이름을 입력한다.
    per_device_train_batch_size=16,
    gradient_accumulation_steps=2,  # 배치 크기가 2배 감소할 때마다 2배씩 증가
    learning_rate=1e-5,
    warmup_steps=1000,
    # max_steps=2,  # epoch 대신 설정
    num_train_epochs=1,     # epoch 수 설정 / max_steps와 이것 중 하나만 설정
    gradient_checkpointing=True,
    fp16=True,
    evaluation_strategy="steps",
    per_device_eval_batch_size=16,
    predict_with_generate=True,
    generation_max_length=225,
    save_steps=1000,
    eval_steps=1000,
    logging_steps=25,
    report_to=["tensorboard"],
    load_best_model_at_end=True,
    metric_for_best_model="cer",  # 한국어의 경우 'wer'보다는 'cer'이 더 적합할 것
    greater_is_better=False,
    push_to_hub=True,
    save_total_limit=5,           # 최대 저장할 모델 수 지정
)