Spaces:
Sleeping
Sleeping
File size: 3,838 Bytes
e3bdc52 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 | # @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.")
|