Spaces:
Sleeping
Sleeping
| # @title 4. Executar Treinamento (Fine-Tuning) | |
| import os | |
| import torch | |
| import librosa | |
| from torch.utils.data import Dataset | |
| from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification, Trainer, TrainingArguments | |
| # Configurações do Modelo | |
| BASE_MODEL = "HyperMoon/wav2vec2-base-960h-finetuned-deepfake" | |
| OUTPUT_DIR = "local_finetuned_model" | |
| # Mapeamento Rígido de Labels para evitar conflitos (0=Real, 1=Fraude) | |
| id2label = {0: "AUTHENTIC", 1: "FAKE"} | |
| label2id = {"AUTHENTIC": 0, "FAKE": 1} | |
| class DeepfakeDataset(Dataset): | |
| def __init__(self, root_dir, processor): | |
| self.files = [] | |
| self.processor = processor | |
| # Carregamento explícito baseado em pastas | |
| for label_name, label_id in label2id.items(): | |
| folder = "real" if label_name == "AUTHENTIC" else "fake" | |
| path = os.path.join(root_dir, folder) | |
| if os.path.exists(path): | |
| print(f"Carregando audios de: {folder}...") | |
| for f in os.listdir(path): | |
| if f.lower().endswith(('.wav', '.mp3', '.flac')): | |
| self.files.append({"path": os.path.join(path, f), "label": label_id}) | |
| else: | |
| print(f"AVISO: Pasta {folder} não encontrada em {root_dir}") | |
| def __len__(self): return len(self.files) | |
| def __getitem__(self, idx): | |
| item = self.files[idx] | |
| try: | |
| speech, _ = librosa.load(item["path"], sr=16000) | |
| inputs = self.processor(speech, sampling_rate=16000, return_tensors="pt", padding="max_length", max_length=160000, truncation=True) | |
| return {"input_values": inputs.input_values[0], "labels": torch.tensor(item["label"])} | |
| except Exception as e: | |
| print(f"Erro ao processar {item['path']}: {e}") | |
| # Retorna o primeiro item como fallback para não quebrar o loop do Trainer | |
| return self.__getitem__(0) | |
| print("Inicializando Processador e Modelo...") | |
| try: | |
| processor = Wav2Vec2FeatureExtractor.from_pretrained(BASE_MODEL) | |
| # Adicionado id2label e label2id aqui para garantir consistência | |
| model = Wav2Vec2ForSequenceClassification.from_pretrained( | |
| BASE_MODEL, | |
| num_labels=2, | |
| id2label=id2label, | |
| label2id=label2id, | |
| ignore_mismatched_sizes=True | |
| ) | |
| # Congelar base para focar no aprendizado das novas fraudes (Lógica Robusta) | |
| if hasattr(model, 'wav2vec2'): | |
| for param in model.wav2vec2.parameters(): | |
| param.requires_grad = False | |
| print("Modelo carregado e camadas base congeladas com sucesso!") | |
| # Dataset (Aponte para a pasta onde você subiu os áudios no Colab) | |
| # Ex: /content/dataset_treino | |
| dataset_path = "/content/dataset" | |
| train_data = DeepfakeDataset(dataset_path, processor) | |
| if len(train_data) == 0: | |
| print("ERRO: Nenhum dado encontrado. Verifique se as pastas 'real' e 'fake' existem dentro do caminho especificado.") | |
| else: | |
| training_args = TrainingArguments( | |
| output_dir=OUTPUT_DIR, | |
| num_train_epochs=3, | |
| per_device_train_batch_size=2, | |
| gradient_accumulation_steps=4, | |
| save_steps=50, | |
| logging_steps=10, | |
| learning_rate=2e-5, | |
| remove_unused_columns=False | |
| ) | |
| trainer = Trainer(model=model, args=training_args, train_dataset=train_data) | |
| print("Iniciando Treinamento...") | |
| trainer.train() | |
| # Salva o resultado final | |
| model.save_pretrained(OUTPUT_DIR) | |
| processor.save_pretrained(OUTPUT_DIR) | |
| print(f"Sucesso! Modelo salvo em: {OUTPUT_DIR}") | |
| except Exception as e: | |
| print(f"ERRO CRÍTICO: {e}") | |
| print("DICA: Se o erro for de conexão, tente rodar a célula novamente. O Hugging Face pode falhar ocasionalmente no download.") | |