maxseats/SungBeom-whisper-small-ko-set1

Model Name : maxseats/SungBeom-whisper-small-ko-set0

Description

• 파인튜닝 데이터셋 : maxseats/aihub-464-preprocessed-680GB-set-1

설명

• AI hub의 주요 영역별 회의 음성 데이터셋을 학습 중이에요.

• 680GB 중 첫번째 데이터(10GB)를 파인튜닝한 모델을 불러와서, 두번째 데이터를 학습한 모델입니다.

• 링크 : https://huggingface.co/datasets/maxseats/aihub-464-preprocessed-680GB-set-0, https://huggingface.co/datasets/maxseats/aihub-464-preprocessed-680GB-set-1

• 다음 코드를 통해 작성했어요.

from datasets import load_dataset
import torch
from dataclasses import dataclass
from typing import Any, Dict, List, Union
import evaluate
from transformers import WhisperTokenizer, WhisperFeatureExtractor, WhisperProcessor, WhisperForConditionalGeneration, Seq2SeqTrainingArguments, Seq2SeqTrainer
import mlflow
from mlflow.tracking.client import MlflowClient
import subprocess
from huggingface_hub import create_repo, Repository
import os
import shutil
import math # 임시 테스트용
model_dir = "./tmpp" # 수정 X


#########################################################################################################################################
################################################### 사용자 설정 변수 #####################################################################
#########################################################################################################################################

model_description = """
- 파인튜닝 데이터셋 : maxseats/aihub-464-preprocessed-680GB-set-1

# 설명
- AI hub의 주요 영역별 회의 음성 데이터셋을 학습 중이에요.
- 680GB 중 첫번째 데이터(10GB)를 파인튜닝한 모델을 불러와서, 두번째 데이터를 학습한 모델입니다.
- 링크 : https://huggingface.co/datasets/maxseats/aihub-464-preprocessed-680GB-set-0, https://huggingface.co/datasets/maxseats/aihub-464-preprocessed-680GB-set-1
"""

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

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=500,
    # 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,           # 최대 저장할 모델 수 지정
)

#########################################################################################################################################
################################################### 사용자 설정 변수 #####################################################################
#########################################################################################################################################


@dataclass
class DataCollatorSpeechSeq2SeqWithPadding:
    processor: Any

    def __call__(self, features: List[Dict[str, Union[List[int], torch.Tensor]]]) -> Dict[str, torch.Tensor]:
        # 인풋 데이터와 라벨 데이터의 길이가 다르며, 따라서 서로 다른 패딩 방법이 적용되어야 한다. 그러므로 두 데이터를 분리해야 한다.
        # 먼저 오디오 인풋 데이터를 간단히 토치 텐서로 반환하는 작업을 수행한다.
        input_features = [{"input_features": feature["input_features"]} for feature in features]
        batch = self.processor.feature_extractor.pad(input_features, return_tensors="pt")

        # Tokenize된 레이블 시퀀스를 가져온다.
        label_features = [{"input_ids": feature["labels"]} for feature in features]
        # 레이블 시퀀스에 대해 최대 길이만큼 패딩 작업을 실시한다.
        labels_batch = self.processor.tokenizer.pad(label_features, return_tensors="pt")

        # 패딩 토큰을 -100으로 치환하여 loss 계산 과정에서 무시되도록 한다.
        labels = labels_batch["input_ids"].masked_fill(labels_batch.attention_mask.ne(1), -100)

        # 이전 토크나이즈 과정에서 bos 토큰이 추가되었다면 bos 토큰을 잘라낸다.
        # 해당 토큰은 이후 언제든 추가할 수 있다.
        if (labels[:, 0] == self.processor.tokenizer.bos_token_id).all().cpu().item():
            labels = labels[:, 1:]

        batch["labels"] = labels

        return batch


def compute_metrics(pred):
    pred_ids = pred.predictions
    label_ids = pred.label_ids

    # pad_token을 -100으로 치환
    label_ids[label_ids == -100] = tokenizer.pad_token_id

    # metrics 계산 시 special token들을 빼고 계산하도록 설정
    pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
    label_str = tokenizer.batch_decode(label_ids, skip_special_tokens=True)

    cer = 100 * metric.compute(predictions=pred_str, references=label_str)

    return {"cer": cer}


# model_dir, ./repo 초기화
if os.path.exists(model_dir):
    shutil.rmtree(model_dir)
os.makedirs(model_dir)

if os.path.exists('./repo'):
    shutil.rmtree('./repo')
os.makedirs('./repo')

# 파인튜닝을 진행하고자 하는 모델의 processor, tokenizer, feature extractor, model 로드
processor = WhisperProcessor.from_pretrained(model_name, language="Korean", task="transcribe")
tokenizer = WhisperTokenizer.from_pretrained(model_name, language="Korean", task="transcribe")
feature_extractor = WhisperFeatureExtractor.from_pretrained(model_name)
model = WhisperForConditionalGeneration.from_pretrained(model_name)

data_collator = DataCollatorSpeechSeq2SeqWithPadding(processor=processor)
metric = evaluate.load('cer')
model.config.forced_decoder_ids = None
model.config.suppress_tokens = []


# Hub로부터 "모든 전처리가 완료된" 데이터셋을 로드(이게 진짜 오래걸려요.)
preprocessed_dataset = load_dataset(dataset_name, cache_dir=CACHE_DIR)

# 30%까지의 valid 데이터셋 선택(코드 작동 테스트를 위함)
if is_test:
    preprocessed_dataset["valid"] = preprocessed_dataset["valid"].select(range(math.ceil(len(preprocessed_dataset) * 0.3)))

# training_args 객체를 JSON 형식으로 변환
training_args_dict = training_args.to_dict()

# MLflow UI 관리 폴더 지정
mlflow.set_tracking_uri("sqlite:////content/drive/MyDrive/STT_test/mlflow.db")

# MLflow 실험 이름을 모델 이름으로 설정
experiment_name = model_name
existing_experiment = mlflow.get_experiment_by_name(experiment_name)

if existing_experiment is not None:
    experiment_id = existing_experiment.experiment_id
else:
    experiment_id = mlflow.create_experiment(experiment_name)


model_version = 1  # 로깅 하려는 모델 버전(이미 존재하면, 자동 할당)

# MLflow 로깅
with mlflow.start_run(experiment_id=experiment_id, description=model_description):
    # training_args 로깅
    for key, value in training_args_dict.items():
        mlflow.log_param(key, value)


    mlflow.set_tag("Dataset", dataset_name) # 데이터셋 로깅

    trainer = Seq2SeqTrainer(
        args=training_args,
        model=model,
        train_dataset=preprocessed_dataset["train"],
        eval_dataset=preprocessed_dataset["valid"],  # or "test"
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        tokenizer=processor.feature_extractor,
    )

    trainer.train()
    trainer.save_model(model_dir)  # 학습 후 모델 저장

    # Metric 로깅
    metrics = trainer.evaluate()
    for metric_name, metric_value in metrics.items():
        mlflow.log_metric(metric_name, metric_value)

    # MLflow 모델 레지스터
    model_uri = "runs:/{run_id}/{artifact_path}".format(run_id=mlflow.active_run().info.run_id, artifact_path=model_dir)

    # 이 값 이용해서 허깅페이스 모델 이름 설정 예정
    model_details = mlflow.register_model(model_uri=model_uri, name=model_name.replace('/', '-'))   # 모델 이름에 '/'를 '-'로 대체

    # 모델 Description
    client = MlflowClient()
    client.update_model_version(name=model_details.name, version=model_details.version, description=model_description)
    model_version = model_details.version   # 버전 정보 허깅페이스 업로드 시 사용



## 허깅페이스 로그인
while True:

    if token =="exit":
        break

    try:
        result = subprocess.run(["huggingface-cli", "login", "--token", token])
        if result.returncode != 0:
            raise Exception()
        break
    except Exception as e:
        token = input("Please enter your Hugging Face API token: ")


os.environ["HUGGINGFACE_HUB_TOKEN"] = token

# 리포지토리 이름 설정
repo_name = "maxseats/" + model_name.replace('/', '-') + '-' + str(model_version)  # 허깅페이스 레포지토리 이름 설정

# 리포지토리 생성
create_repo(repo_name, exist_ok=True, token=token)



# 리포지토리 클론
repo = Repository(local_dir='./repo', clone_from=f"{repo_name}", use_auth_token=token)


# model_dir 필요한 파일 복사
max_depth = 1  # 순회할 최대 깊이

for root, dirs, files in os.walk(model_dir):
    depth = root.count(os.sep) - model_dir.count(os.sep)
    if depth < max_depth:
        for file in files:
            # 파일 경로 생성
            source_file = os.path.join(root, file)
            # 대상 폴더에 복사
            shutil.copy(source_file, './repo')


# 토크나이저 다운로드 및 로컬 디렉토리에 저장
tokenizer.save_pretrained('./repo')


readme = f"""
---
language: ko
tags:
- whisper
- speech-recognition
datasets:
- {dataset_name}
metrics:
- cer
---
# Model Name : {model_name}
# Description
{model_description}
"""


# 모델 카드 및 기타 메타데이터 파일 작성
with open("./repo/README.md", "w") as f:
    f.write(readme)

# 파일 커밋 푸시
repo.push_to_hub(commit_message="Initial commit")

# 폴더와 하위 내용 삭제
shutil.rmtree(model_dir)
shutil.rmtree('./repo')