train.py

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()