main.py

import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader, random_split
import pandas as pd
import numpy as np
from transformers import (
    AutoTokenizer, AutoModelForSequenceClassification,
    TrainingArguments, Trainer, DataCollatorWithPadding
)
from sklearn.metrics import accuracy_score, classification_report, confusion_matrix
import matplotlib.pyplot as plt
import seaborn as sns
from datasets import Dataset as HFDataset
import warnings
import os
from multiprocessing import cpu_count

warnings.filterwarnings('ignore')


class UsernameDataset(Dataset):
    def __init__(self, real_csv=None, synthetic_csv=None, combined_csv=None, tokenizer=None, max_length=128):
        self.tokenizer = tokenizer
        self.max_length = max_length

        if combined_csv:
            print("Loading combined CSV...")
            self.data = pd.read_csv(combined_csv)
            if 'username' not in self.data.columns or 'label' not in self.data.columns:
                raise ValueError("Combined CSV must contain 'username' and 'label' columns")
            self.usernames = self.data['username'].tolist()
            self.labels = self.data['label'].tolist()
        else:
            if not (real_csv and synthetic_csv):
                raise ValueError("Must provide either combined_csv or both real_csv and synthetic_csv")

            print("Loading real usernames CSV...")
            real_data = pd.read_csv(real_csv)
            print("Loading synthetic usernames CSV...")
            synthetic_data = pd.read_csv(synthetic_csv)

            username_col = 'username' if 'username' in real_data.columns else real_data.columns[0]

            real_usernames = real_data[username_col].tolist()
            synthetic_usernames = synthetic_data[username_col].tolist()

            self.usernames = real_usernames + synthetic_usernames
            self.labels = [0] * len(real_usernames) + [1] * len(synthetic_usernames)  # 0 = real, 1 = AI generated

        print(f"Loaded {len(self.usernames)} usernames")
        print(f"Real usernames: {self.labels.count(0)}, Generated usernames: {self.labels.count(1)}")

    def __len__(self):
        return len(self.usernames)

    def __getitem__(self, idx):
        username = str(self.usernames[idx])
        label = self.labels[idx]
        return {'text': username, 'labels': label}


def compute_metrics(eval_pred):
    preds, labels = eval_pred
    preds = np.argmax(preds, axis=1)
    acc = accuracy_score(labels, preds)
    return {'accuracy': acc}


def tokenize_function(examples, tokenizer, max_length=64):
    return tokenizer(
        examples['text'],
        truncation=True,
        padding='max_length',
        max_length=max_length,
        return_tensors='pt'
    )


def train_llm_model(model, train_dataset, val_dataset, tokenizer, output_dir='./username_classifier_llm',
                    num_epochs=3, batch_size=32, learning_rate=2e-5, gradient_accumulation_steps=1, max_length=64):
    # this is where the actual training happens,  teaching the model to recognize patterns
    print("Converting datasets to HuggingFace format...")
    train_texts = [train_dataset[i]['text'] for i in range(len(train_dataset))]
    train_labels = [train_dataset[i]['labels'] for i in range(len(train_dataset))]

    val_texts = [val_dataset[i]['text'] for i in range(len(val_dataset))]
    val_labels = [val_dataset[i]['labels'] for i in range(len(val_dataset))]

    train_hf_dataset = HFDataset.from_dict({
        'text': train_texts,
        'labels': train_labels
    })

    val_hf_dataset = HFDataset.from_dict({
        'text': val_texts,
        'labels': val_labels
    })

    from functools import partial
    tokenize_fn = partial(tokenize_function, tokenizer=tokenizer, max_length=max_length)

    # tokenizing=converting text to numbers the model can understand
    print("Tokenizing training dataset...")
    num_proc = min(cpu_count(), 4)

    train_hf_dataset = train_hf_dataset.map(
        tokenize_fn,
        batched=True,
        batch_size=200,
        num_proc=num_proc,
        remove_columns=['text'],
        desc="Tokenizing train"
    )

    print("Tokenizing validation dataset...")
    val_hf_dataset = val_hf_dataset.map(
        tokenize_fn,
        batched=True,
        batch_size=200,
        num_proc=num_proc,
        remove_columns=['text'],
        desc="Tokenizing validation"
    )

    data_collator = DataCollatorWithPadding(tokenizer=tokenizer)

    total_steps = len(train_hf_dataset) // (batch_size * gradient_accumulation_steps) * num_epochs
    warmup_steps = int(0.1 * total_steps)  # gradually increasing learning rate helps the model learn better

    # this took me forever to figure out please dont touch thanks
    training_args = TrainingArguments(
        output_dir=output_dir,
        num_train_epochs=num_epochs,
        per_device_train_batch_size=batch_size,
        per_device_eval_batch_size=batch_size,
        gradient_accumulation_steps=gradient_accumulation_steps,
        learning_rate=learning_rate,
        weight_decay=0.01,
        logging_dir='./logs',
        logging_steps=10,
        eval_strategy='steps',
        eval_steps=1000,
        save_strategy='steps',
        save_steps=2000,
        load_best_model_at_end=True,
        metric_for_best_model='accuracy',
        greater_is_better=True,
        warmup_steps=warmup_steps,
        save_total_limit=2,
        report_to=None,
        fp16=True,  # mixed precision is  faster training on GPU
        dataloader_pin_memory=True,
        dataloader_num_workers=min(4, cpu_count()//2),
        group_by_length=True,
        gradient_checkpointing=True,  # saves memory but slows down training a bit
        optim='adamw_torch_fused',
        disable_tqdm=False,
        prediction_loss_only=False,
        logging_first_step=True,
        log_level='info',
        max_steps=10000,
        max_grad_norm=1.0,
    )

    from transformers import TrainerCallback
    import time

    class ProgressCallback(TrainerCallback):
        def __init__(self):
            self.start_time = time.time()
            self.step_times = []

        def on_step_end(self, args, state, control, **kwargs):
            current_time = time.time()
            if len(self.step_times) > 0:
                step_time = current_time - self.step_times[-1]
                if state.global_step % 10 == 0:
                    elapsed = current_time - self.start_time
                    print(f"Step {state.global_step}: {step_time:.2f}s/step, Total: {elapsed/60:.1f}min")
            self.step_times.append(current_time)

        def on_epoch_end(self, args, state, control, **kwargs):
            elapsed = time.time() - self.start_time
            print(f"Epoch {state.epoch} completed in {elapsed/60:.1f} minutes")

    trainer = Trainer(
        model=model,
        args=training_args,
        train_dataset=train_hf_dataset,
        eval_dataset=val_hf_dataset,
        tokenizer=tokenizer,
        data_collator=data_collator,
        compute_metrics=compute_metrics,
        callbacks=[ProgressCallback()],
    )

    print("Starting training...")
    trainer.train()

    trainer.save_model()
    tokenizer.save_pretrained(output_dir)

    return trainer


def evaluate_llm_model(model, test_dataset, tokenizer, device='cpu', batch_size=64):
    model.eval()
    all_preds = []
    all_labels = []

    test_texts = [test_dataset[i]['text'] for i in range(len(test_dataset))]
    test_labels = [test_dataset[i]['labels'] for i in range(len(test_dataset))]

    model.to(device)

    batch_size = min(batch_size, len(test_texts))

    with torch.no_grad():
        for i in range(0, len(test_texts), batch_size):
            batch_texts = test_texts[i:i+batch_size]
            batch_labels = test_labels[i:i+batch_size]

            inputs = tokenizer(
                batch_texts,
                return_tensors='pt',
                truncation=True,
                padding=True,
                max_length=128
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}

            outputs = model(**inputs)
            preds = torch.argmax(outputs.logits, dim=-1).cpu().numpy()

            all_preds.extend(preds)
            all_labels.extend(batch_labels)

    acc = accuracy_score(all_labels, all_preds)
    report = classification_report(all_labels, all_preds, target_names=['Real', 'Generated'])
    cm = confusion_matrix(all_labels, all_preds)

    return acc, report, cm, all_preds, all_labels


def main():
    # paths to the data files
    REAL_CSV = 'real.csv'
    SYNTHETIC_CSV = 'synthetic.csv'

    MODEL_NAME = 'distilbert-base-uncased'  # using distilbert because it's smaller/faster than full BERT (im not made of money)

    BATCH_SIZE = 64
    GRADIENT_ACCUMULATION_STEPS = 1
    LEARNING_RATE = 3e-5
    NUM_EPOCHS = 1  # might need more epochs but starting with 1 because dataset is huge
    MAX_LENGTH = 64

    TOTAL_SAMPLES = 1500000
    TRAIN_SAMPLES = 1200000
    VAL_SAMPLES = 200000
    TEST_SAMPLES = 100000

    USE_SUBSET = False  # set to True  to test on smaller data first
    SUBSET_SIZE = 50000

    if torch.cuda.is_available():
        device = torch.device('cuda')
        print(f'Using GPU: {torch.cuda.get_device_name(0)}')
        print(f'GPU Memory: {torch.cuda.get_device_properties(0).total_memory / 1024**3:.1f} GB')

        # these settings help GPU run faster
        torch.backends.cudnn.benchmark = True
        torch.backends.cudnn.enabled = True

        if torch.cuda.get_device_properties(0).total_memory > 8 * 1024**3:
            BATCH_SIZE = 64
            GRADIENT_ACCUMULATION_STEPS = 1
    else:
        device = torch.device('cpu')
        print('Using CPU')
        BATCH_SIZE = 16
        GRADIENT_ACCUMULATION_STEPS = 1

    print(f"Loading tokenizer and model: {MODEL_NAME}")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token

    # num_labels=2 because we're doing binary classification (real vs fake/synthetic)
    model = AutoModelForSequenceClassification.from_pretrained(
        MODEL_NAME,
        num_labels=2,
        ignore_mismatched_sizes=True
    )

    model.resize_token_embeddings(len(tokenizer))
    model.to(device)

    print("Loading dataset...")
    full_dataset = UsernameDataset(
        real_csv=REAL_CSV,
        synthetic_csv=SYNTHETIC_CSV,
        tokenizer=None,
        max_length=MAX_LENGTH
    )

    if USE_SUBSET:
        print(f"Using subset of {SUBSET_SIZE} samples for testing...")
        dataset_size = min(SUBSET_SIZE, len(full_dataset))

        # keeping the same ratio of real to fake usernames
        total_real = full_dataset.labels.count(0)
        total_synthetic = full_dataset.labels.count(1)

        real_ratio = total_real / (total_real + total_synthetic)
        synthetic_ratio = total_synthetic / (total_real + total_synthetic)

        subset_real = int(dataset_size * real_ratio)
        subset_synthetic = dataset_size - subset_real

        real_indices = [i for i, label in enumerate(full_dataset.labels) if label == 0]
        synthetic_indices = [i for i, label in enumerate(full_dataset.labels) if label == 1]

        selected_real = np.random.choice(real_indices, min(subset_real, len(real_indices)), replace=False)
        selected_synthetic = np.random.choice(synthetic_indices, min(subset_synthetic, len(synthetic_indices)), replace=False)

        subset_indices = np.concatenate([selected_real, selected_synthetic])
        np.random.shuffle(subset_indices)

        subset_usernames = [full_dataset.usernames[i] for i in subset_indices]
        subset_labels = [full_dataset.labels[i] for i in subset_indices]

        train_size = int(0.8 * len(subset_usernames))
        val_size = int(0.15 * len(subset_usernames))
        test_size = len(subset_usernames) - train_size - val_size

        print(f"Subset splits - Train: {train_size}, Val: {val_size}, Test: {test_size}")

    else:
        print(f"Using full dataset with target {TOTAL_SAMPLES} samples...")

        if len(full_dataset) > TOTAL_SAMPLES:
            total_real = full_dataset.labels.count(0)
            total_synthetic = full_dataset.labels.count(1)

            real_ratio = total_real / (total_real + total_synthetic)
            synthetic_ratio = total_synthetic / (total_real + total_synthetic)

            target_real = int(TOTAL_SAMPLES * real_ratio)
            target_synthetic = TOTAL_SAMPLES - target_real

            real_indices = [i for i, label in enumerate(full_dataset.labels) if label == 0]
            synthetic_indices = [i for i, label in enumerate(full_dataset.labels) if label == 1]

            selected_real = np.random.choice(real_indices, min(target_real, len(real_indices)), replace=False)
            selected_synthetic = np.random.choice(synthetic_indices, min(target_synthetic, len(synthetic_indices)), replace=False)

            subset_indices = np.concatenate([selected_real, selected_synthetic])
            np.random.shuffle(subset_indices)

            subset_usernames = [full_dataset.usernames[i] for i in subset_indices]
            subset_labels = [full_dataset.labels[i] for i in subset_indices]

            print(f"Sampled {len(subset_usernames)} samples from {len(full_dataset)} total")
            print(f"Real: {subset_labels.count(0)}, Synthetic: {subset_labels.count(1)}")

        else:
            subset_usernames = full_dataset.usernames
            subset_labels = full_dataset.labels

            actual_size = len(subset_usernames)
            TRAIN_SAMPLES = int(0.8 * actual_size)
            VAL_SAMPLES = int(0.133 * actual_size)
            TEST_SAMPLES = actual_size - TRAIN_SAMPLES - VAL_SAMPLES

            print(f"Using full dataset of {actual_size} samples")

        train_size = TRAIN_SAMPLES
        val_size = VAL_SAMPLES
        test_size = TEST_SAMPLES

        print(f"Target splits - Train: {train_size}, Val: {val_size}, Test: {test_size}")

    class CustomDataset(Dataset):
        def __init__(self, usernames, labels):
            self.usernames = usernames
            self.labels = labels
        def __len__(self):
            return len(self.usernames)
        def __getitem__(self, idx):
            return {'text': str(self.usernames[idx]), 'labels': self.labels[idx]}

    dataset = CustomDataset(subset_usernames, subset_labels)

    print("Creating custom dataset splits...")

    # randomly shuffling so the model doesn't learn any order patterns
    indices = list(range(len(dataset)))
    np.random.shuffle(indices)

    train_indices = indices[:train_size]
    val_indices = indices[train_size:train_size + val_size]
    test_indices = indices[train_size + val_size:train_size + val_size + test_size]

    class IndexDataset(Dataset):
        def __init__(self, base_dataset, indices):
            self.base_dataset = base_dataset
            self.indices = indices
        def __len__(self):
            return len(self.indices)
        def __getitem__(self, idx):
            return self.base_dataset[self.indices[idx]]

    train_dataset = IndexDataset(dataset, train_indices)
    val_dataset = IndexDataset(dataset, val_indices)
    test_dataset = IndexDataset(dataset, test_indices)

    print(f"Final dataset splits - Train: {len(train_dataset)}, Val: {len(val_dataset)}, Test: {len(test_dataset)}")

    def check_balance(split_dataset, split_name):
        labels = [split_dataset[i]['labels'] for i in range(len(split_dataset))]
        real_count = labels.count(0)
        synthetic_count = labels.count(1)
        print(f"{split_name} - Real: {real_count}, Synthetic: {synthetic_count}, Balance: {real_count/(real_count+synthetic_count):.3f}")

    check_balance(train_dataset, "Train")
    check_balance(val_dataset, "Val")
    check_balance(test_dataset, "Test")

    print(f"Training with {len(train_dataset)} samples...")
    print(f"Effective batch size: {BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS}")
    print(f"Expected steps per epoch: {len(train_dataset) // (BATCH_SIZE * GRADIENT_ACCUMULATION_STEPS)}")

    from transformers import EarlyStoppingCallback

    trainer = train_llm_model(
        model=model,
        train_dataset=train_dataset,
        val_dataset=val_dataset,
        tokenizer=tokenizer,
        num_epochs=NUM_EPOCHS,
        batch_size=BATCH_SIZE,
        learning_rate=LEARNING_RATE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION_STEPS,
        max_length=MAX_LENGTH
    )

    print("\nEvaluating on test set...")
    test_acc, test_report, test_cm, preds, true_labels = evaluate_llm_model(
        model, test_dataset, tokenizer, device, batch_size=BATCH_SIZE*2
    )

    print(f"\nTest Accuracy: {test_acc:.4f}")
    print("\nClassification Report:")
    print(test_report)

    # confusion matrix helps see where the model is making mistakes
    plt.figure(figsize=(8, 6))
    sns.heatmap(test_cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=['Real', 'Generated'], yticklabels=['Real', 'Generated'])
    plt.title('Confusion Matrix')
    plt.ylabel('True Label')
    plt.xlabel('Predicted Label')
    plt.tight_layout()
    plt.show()

    print(f"\nModel and tokenizer saved to './username_classifier_llm'")


def predict_username_llm(model_path, tokenizer_path, username, device='cpu'):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.to(device)
    model.eval()

    inputs = tokenizer(
        username,
        return_tensors='pt',
        truncation=True,
        padding='max_length',
        max_length=128
    )
    inputs = {k: v.to(device) for k, v in inputs.items()}

    with torch.no_grad():
        outputs = model(**inputs)
        probs = torch.softmax(outputs.logits, dim=-1)
        pred_class = torch.argmax(outputs.logits, dim=-1).item()
        conf = probs[0][pred_class].item()

    return pred_class, conf


def predict_batch_usernames(model_path, tokenizer_path, usernames, device='cpu', batch_size=64):
    tokenizer = AutoTokenizer.from_pretrained(tokenizer_path)
    model = AutoModelForSequenceClassification.from_pretrained(model_path)
    model.to(device)
    model.eval()

    preds = []
    confs = []

    with torch.no_grad():
        for i in range(0, len(usernames), batch_size):
            batch_usernames = usernames[i:i+batch_size]

            inputs = tokenizer(
                batch_usernames,
                return_tensors='pt',
                truncation=True,
                padding=True,
                max_length=128
            )
            inputs = {k: v.to(device) for k, v in inputs.items()}

            outputs = model(**inputs)
            probs = torch.softmax(outputs.logits, dim=-1)
            batch_preds = torch.argmax(outputs.logits, dim=-1).cpu().numpy()
            batch_confs = torch.max(probs, dim=-1)[0].cpu().numpy()

            preds.extend(batch_preds)
            confs.extend(batch_confs)

    return preds, confs


def example_predictions():
    model_path = './username_classifier_llm'
    tokenizer_path = './username_classifier_llm'

    test_usernames = [
        'john_doe123',
        'xX_GamerPro_Xx',
        'user_12345',
        'randomuser789',
        'john.smith',
        'generated_user_001'
    ]

    print("Example predictions (batch):")
    try:
        preds, confs = predict_batch_usernames(
            model_path, tokenizer_path, test_usernames
        )

        for username, pred, conf in zip(test_usernames, preds, confs):
            label = 'Generated' if pred == 1 else 'Real'
            print(f"Username: {username:20} | Prediction: {label:10} | Confidence: {conf:.3f}")
    except Exception as e:
        print(f"Error in batch prediction: {e}")


if __name__ == "__main__":
    main()