TTS AI 예제 2

import os
import json
import numpy as np
import librosa
import torch
from transformers import Wav2Vec2ForCTC, Wav2Vec2Processor, Trainer, TrainingArguments
from datasets import Dataset, load_metric
from sklearn.model_selection import train_test_split
import evaluate

# 음성 파일 및 라벨 데이터 폴더 경로 설정
audio_dir = '/content/path_to_audio_folder'  # Colab에서는 content 폴더 내에 경로 설정
label_dir = '/content/path_to_label_folder'  # Colab에서는 content 폴더 내에 경로 설정

# 파일 확장자가 .wav인 파일 목록 가져오기
audio_files = [f for f in os.listdir(audio_dir) if f.endswith('.wav')]

# 파일 확장자가 .json인 라벨 파일 목록 가져오기
label_files = [f for f in os.listdir(label_dir) if f.endswith('.json')]

# 음성 파일을 로드하고 전처리 함수 정의
def load_audio_file(file_path, sr=16000):
    y, sr = librosa.load(file_path, sr=sr)
    return y, sr

# 모든 라벨 파일 로드 및 병합
all_labels = {}
for label_file in label_files:
    with open(os.path.join(label_dir, label_file), 'r', encoding='utf-8') as f:
        labels = json.load(f)
        file_name = labels['파일정보']['FileName']
        all_labels[file_name] = labels

# 데이터 준비
X = []
y = []

for audio_file in audio_files:
    file_path = os.path.join(audio_dir, audio_file)
    if audio_file in all_labels:
        y_audio, sr = load_audio_file(file_path)
        label_data = all_labels[audio_file]
        if "OrgLabelText" in label_data["전사정보"]:
            transcript = label_data["전사정보"]["OrgLabelText"]
            X.append({"array": y_audio, "sampling_rate": sr})
            y.append(transcript)

# X, y 출력 확인
print("X: ", len(X))  # 데이터 개수 확인
print("y: ", len(y))  # 데이터 개수 확인

# train_test_split을 사용하여 데이터 분할
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 데이터셋 준비
train_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_train, y_train)])
test_data = Dataset.from_list([{"audio": X, "transcript": y} for X, y in zip(X_test, y_test)])

# Wav2Vec 2.0 모델 및 프로세서 로드
processor = Wav2Vec2Processor.from_pretrained("facebook/wav2vec2-large-960h")
model = Wav2Vec2ForCTC.from_pretrained("facebook/wav2vec2-large-960h")

# 데이터셋 전처리
def prepare_dataset(batch):
    audio = batch["audio"]
    inputs = processor(audio["array"], sampling_rate=audio["sampling_rate"], return_tensors="pt", padding="max_length", truncation=True, max_length=16000)
    labels = processor(text=[batch["transcript"]], return_tensors="pt", padding="max_length", truncation=True, max_length=128).input_ids
    batch["input_values"] = inputs.input_values.squeeze().numpy()
    batch["labels"] = labels.squeeze().numpy()
    return batch

# 데이터셋 처리
train_dataset = train_data.map(prepare_dataset, remove_columns=["audio", "transcript"])
test_dataset = test_data.map(prepare_dataset, remove_columns=["audio", "transcript"])

# 데이터셋 크기 확인
print(f"Train dataset size: {len(train_dataset)}")
print(f"Test dataset size: {len(test_dataset)}")

# 모델 학습 설정
training_args = TrainingArguments(
    output_dir="./results",
    eval_strategy="epoch",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    save_steps=10,
    save_total_limit=2,
)

# 트레이너 설정
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=test_dataset,
    tokenizer=processor.feature_extractor,
)

# 모델 학습
trainer.train()

# 모델 평가
wer_metric = evaluate.load("wer")

def compute_metrics(pred):
    pred_logits = pred.predictions
    pred_ids = np.argmax(pred_logits, axis=-1)
    pred_str = processor.batch_decode(pred_ids)
    wer = wer_metric.compute(predictions=pred_str, references=pred.label_ids)
    return {"wer": wer}

results = trainer.evaluate()
print(f"Word Error Rate (WER): {results['eval_wer'] * 100:.2f}%")