Update README.md

README.md

@@ -11,10 +11,288 @@ metrics:
 ---
 # Model Name : maxseats/SungBeom-whisper-small-ko-set0
 # Description
+- 파인튜닝 데이터셋 : maxseats/aihub-464-preprocessed-680GB-set-1
 
-- 파인튜닝 데이터셋 : maxseats/aihub-464-preprocessed-680GB-set-0
+# 설명
+- 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
 
 # 설명
-- 주요 영역별 회의 음성 데이터셋 680GB 중 첫번째 데이터(10GB)를 파인튜닝한 모델입니다.
-- 링크 : https://huggingface.co/datasets/maxseats/aihub-464-preprocessed-680GB-set-0
+- 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')
+```
+