import os import torch import glob from datasets import load_dataset, Audio, Dataset, ClassLabel from transformers import ( AutoFeatureExtractor, AutoModelForAudioClassification, TrainingArguments, Trainer ) import numpy as np import evaluate import random # --- Configurações --- # Modelo base do Hugging Face. XLSR-53 é uma ótima escolha multilíngue. MODEL_NAME = "lgris/w2v_podcasts_base_400k_pt" # Pasta onde os dados pré-processados foram salvos DATASET_PATH = "./dataset_preparado/" # Pasta para salvar o modelo treinado e os logs OUTPUT_DIR = "./portuguese-accent-classifier" # Hiperparâmetros de treinamento NUM_TRAIN_EPOCHS = 50 BATCH_SIZE = 32 # Reduza se tiver problemas de memória na GPU def train_model(): """ Carrega o dataset, pré-processa os dados e faz o fine-tuning do modelo para classificação de sotaques. """ print("Carregando dataset a partir das pastas...") # Carregar todos os arquivos manualmente para contornar limitação do audiofolder pt_br_files = glob.glob(os.path.join(DATASET_PATH, "pt_br", "*.wav")) pt_pt_files = glob.glob(os.path.join(DATASET_PATH, "pt_pt", "*.wav")) print(f"Arquivos pt_br encontrados: {len(pt_br_files)}") print(f"Arquivos pt_pt encontrados: {len(pt_pt_files)}") # Criar listas de arquivos e labels all_files = [] all_labels = [] # Adicionar arquivos pt_br com label 0 for file_path in pt_br_files: all_files.append(file_path) all_labels.append(0) # Adicionar arquivos pt_pt com label 1 for file_path in pt_pt_files: all_files.append(file_path) all_labels.append(1) print(f"Total de arquivos carregados: {len(all_files)}") # Criar dataset customizado data_dict = { "audio": all_files, "label": all_labels } dataset = Dataset.from_dict(data_dict) dataset = dataset.cast_column("audio", Audio()) # Configurar labels como ClassLabel para permitir estratificação labels = ["pt_br", "pt_pt"] dataset = dataset.cast_column("label", ClassLabel(names=labels)) print(f"Dataset criado com {len(dataset)} exemplos") # Dividir em treino e teste dataset = dataset.train_test_split(test_size=0.1, shuffle=True, stratify_by_column="label") print("Dataset carregado e dividido:") print(dataset) # Criar mapeamento de labels label2id, id2label = {}, {} for i, label in enumerate(labels): label2id[label] = str(i) id2label[str(i)] = label print(f"Mapeamento de labels: {id2label}") # Carregar o Feature Extractor feature_extractor = AutoFeatureExtractor.from_pretrained(MODEL_NAME) target_sampling_rate = feature_extractor.sampling_rate target_length = int(target_sampling_rate * 5.0) # 5 segundos em samples def preprocess_function(examples): # O dataset já foi resampleado, mas garantimos aqui # A função Audio() do 'datasets' carrega o áudio audio_arrays = [x["array"] for x in examples["audio"]] # Processar cada áudio individualmente para truncamento/padding aleatório processed_audios = [] for audio_array in audio_arrays: audio_length = len(audio_array) if audio_length > target_length: # Áudio maior que 5s: truncar aleatoriamente start_idx = random.randint(0, audio_length - target_length) processed_audio = audio_array[start_idx:start_idx + target_length] else: # Áudio menor que 5s: adicionar silêncio aleatoriamente padding_needed = target_length - audio_length # Distribuir o padding aleatoriamente entre início e fim left_padding = random.randint(0, padding_needed) right_padding = padding_needed - left_padding # Criar arrays de silêncio (zeros) left_silence = np.zeros(left_padding, dtype=audio_array.dtype) right_silence = np.zeros(right_padding, dtype=audio_array.dtype) # Concatenar: silêncio_esquerdo + áudio + silêncio_direito processed_audio = np.concatenate([left_silence, audio_array, right_silence]) processed_audios.append(processed_audio) # Usar o feature_extractor com os áudios já processados inputs = feature_extractor( processed_audios, sampling_rate=target_sampling_rate, padding=False, # Não precisamos de padding adicional truncation=False # Não precisamos de truncamento adicional ) return inputs # Aplicar o pré-processamento ao dataset print("Aplicando pré-processamento...") encoded_dataset = dataset.map(preprocess_function, remove_columns="audio", batched=True) # Carregar o Modelo num_labels = len(labels) model = AutoModelForAudioClassification.from_pretrained( MODEL_NAME, num_labels=num_labels, label2id={k: int(v) for k, v in label2id.items()}, # Convertendo keys de str para int id2label=id2label, ignore_mismatched_sizes=True # Permite substituir a 'cabeça' de classificação ) # Métrica de avaliação accuracy = evaluate.load("accuracy") def compute_metrics(eval_pred): predictions = np.argmax(eval_pred.predictions, axis=1) return accuracy.compute(predictions=predictions, references=eval_pred.label_ids) # Argumentos de Treinamento training_args = TrainingArguments( output_dir=OUTPUT_DIR, eval_strategy="epoch", # Adicionado para compatibilidade com load_best_model_at_end save_strategy="epoch", learning_rate=3e-5, per_device_train_batch_size=BATCH_SIZE, per_device_eval_batch_size=BATCH_SIZE, num_train_epochs=NUM_TRAIN_EPOCHS, weight_decay=0.01, logging_steps=10, load_best_model_at_end=True, metric_for_best_model="accuracy", report_to="tensorboard", push_to_hub=False, ) # Inicializar o Trainer trainer = Trainer( model=model, args=training_args, train_dataset=encoded_dataset["train"], eval_dataset=encoded_dataset["test"], tokenizer=feature_extractor, compute_metrics=compute_metrics, ) # Iniciar o Treinamento print("\n--- Iniciando o Fine-tuning ---") trainer.train() print("--- Treinamento Concluído ---\n") # Salvar o modelo final final_model_path = os.path.join(OUTPUT_DIR, "final_model") trainer.save_model(final_model_path) print(f"Modelo final salvo em: {final_model_path}") if __name__ == '__main__': # Garante que o treinamento use GPU se disponível if not torch.cuda.is_available(): print("Aviso: Nenhuma GPU encontrada. O treinamento será muito lento.") train_model()