train_wav2vec2.py

import os
import random
import librosa
import numpy as np
import torch

from datasets import Dataset
from transformers import (
    Wav2Vec2FeatureExtractor,
    Wav2Vec2ForSequenceClassification,
    TrainingArguments,
    Trainer
)

# =========================
# CONFIG
# =========================
BASE = "/home/aids/Downloads/HCL-DATA"

HUMAN_DIR = os.path.join(BASE, "human")
AI_DIR = os.path.join(BASE, "ai_generated")

LANGS = ["tamil","english","hindi","malayalam","telugu"]

SAMPLES_PER_LANG_PER_CLASS = 1500
MAX_SEC = 7
MODEL_NAME = "facebook/wav2vec2-large-xlsr-53"

# =========================
# AUDIO AUGMENTATIONS
# =========================
def augment_audio(y):
    if random.random() < 0.5:
        # random gain
        y = y * random.uniform(0.8, 1.2)

    if random.random() < 0.3:
        # noise injection
        noise = np.random.randn(len(y)) * 0.003
        y = y + noise

    if random.random() < 0.3:
        # pitch shift
        y = librosa.effects.pitch_shift(y, sr=16000, n_steps=random.uniform(-2,2))

    if random.random() < 0.3:
        # time stretch
        rate = random.uniform(0.9, 1.1)
        y = librosa.effects.time_stretch(y, rate)

    return y

# =========================
# COLLECT BALANCED DATA
# =========================
data = []

for lang in LANGS:
    print("Collecting:", lang)

    human_files = os.listdir(f"{HUMAN_DIR}/{lang}/clips")
    ai_files = os.listdir(f"{AI_DIR}/{lang}")

    random.shuffle(human_files)
    random.shuffle(ai_files)

    human_files = human_files[:SAMPLES_PER_LANG_PER_CLASS]
    ai_files = ai_files[:SAMPLES_PER_LANG_PER_CLASS]

    for f in human_files:
        data.append({"path": f"{HUMAN_DIR}/{lang}/clips/{f}", "label": 0})

    for f in ai_files:
        data.append({"path": f"{AI_DIR}/{lang}/{f}", "label": 1})

print("Total samples:", len(data))
dataset = Dataset.from_list(data)

# =========================
# FEATURE EXTRACTOR
# =========================
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_NAME)

def preprocess(example):
    try:
        speech, sr = librosa.load(example["path"], sr=16000)

        # crop
        speech = speech[:16000 * MAX_SEC]

        # augment only for training
        speech = augment_audio(speech)

        speech = speech.astype(np.float32)

        inputs = feature_extractor(
            speech,
            sampling_rate=16000,
            return_tensors="np"
        )

        return {
            "input_values": inputs["input_values"][0],
            "label": example["label"]
        }

    except:
        return {"input_values": None, "label": example["label"]}

dataset = dataset.map(preprocess, num_proc=4)
dataset = dataset.filter(lambda x: x["input_values"] is not None)

# =========================
# SPLIT
# =========================
split = dataset.train_test_split(test_size=0.1)
train_ds = split["train"]
test_ds = split["test"]

# =========================
# MODEL
# =========================
model = Wav2Vec2ForSequenceClassification.from_pretrained(
    MODEL_NAME,
    num_labels=2
)

# Freeze lower layers (critical)
model.freeze_feature_encoder()

# =========================
# TRAINING SETTINGS
# =========================
training_args = TrainingArguments(
    output_dir="./FINAL_ai_voice_detector",

    per_device_train_batch_size=1,
    per_device_eval_batch_size=1,

    gradient_accumulation_steps=8,

    evaluation_strategy="epoch",
    save_strategy="epoch",

    num_train_epochs=7,
    learning_rate=2e-6,
    warmup_steps=500,

    fp16=True,
    logging_steps=20,
    save_total_limit=2,

    load_best_model_at_end=True,
    metric_for_best_model="eval_loss"
)

# =========================
# TRAINER
# =========================
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_ds,
    eval_dataset=test_ds,
    tokenizer=feature_extractor
)

# =========================
# TRAIN
# =========================
trainer.train()

# =========================
# UNFREEZE + FINE POLISH
# =========================
print("Fine tuning full network...")

for param in model.parameters():
    param.requires_grad = True

training_args.num_train_epochs = 2
training_args.learning_rate = 1e-6

trainer.train()

# =========================
# SAVE FINAL MODEL
# =========================
trainer.save_model("./FINAL_ai_voice_detector")
feature_extractor.save_pretrained("./FINAL_ai_voice_detector")

print("ULTIMATE MODEL READY")