ank52/logic_stream

1import os2from datasets import load_dataset3from transformers import (4    AutoTokenizer, 5    AutoModelForSequenceClassification, 6    TrainingArguments, 7    Trainer8)9import evaluate10import numpy as np11 12def compute_metrics(eval_pred):13    metric = evaluate.load("accuracy")14    logits, labels = eval_pred15    predictions = np.argmax(logits, axis=-1)16    return metric.compute(predictions=predictions, references=labels)17 18def train_content_moderation(sample_size=None):19    print("🚀 Starting Content Moderation Model Training Pipeline...")20    21    # 1. Load a massive "Big Data" dataset (Toxic Comments)22    # Using 'tweets_hate_speech_detection' as a public dataset example23    print("⬇️ Downloading Dataset...")24    dataset = load_dataset("tweet_eval", "hate")25    26    if sample_size:27        print(f"⚠️ Demo Mode: Only training on {sample_size} samples for speed.")28        train_dataset = dataset["train"].select(range(sample_size))29        eval_dataset = dataset["validation"].select(range(min(sample_size, len(dataset["validation"]))))30    else:31        print("🧠 Big Data Mode: Training on full dataset!")32        train_dataset = dataset["train"]33        eval_dataset = dataset["validation"]34 35    # 2. Tokenization36    model_name = "distilbert-base-uncased"37    print(f"🔨 Initializing Tokenizer: {model_name}")38    tokenizer = AutoTokenizer.from_pretrained(model_name)39 40    def tokenize_function(examples):41        return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=128)42 43    tokenized_train = train_dataset.map(tokenize_function, batched=True)44    tokenized_eval = eval_dataset.map(tokenize_function, batched=True)45 46    # 3. Model Initialization47    # Label 0: Safe, Label 1: Hate/Toxic48    print("🤖 Loading DistilBERT architecture...")49    model = AutoModelForSequenceClassification.from_pretrained(model_name, num_labels=2)50 51    # 4. Training Configuration52    training_args = TrainingArguments(53        output_dir="./results_moderation",54        evaluation_strategy="epoch",55        learning_rate=2e-5,56        per_device_train_batch_size=16,57        per_device_eval_batch_size=16,58        num_train_epochs=3 if not sample_size else 1, # Fast epoch for demo59        weight_decay=0.01,60        save_strategy="epoch",61        load_best_model_at_end=True,62    )63 64    trainer = Trainer(65        model=model,66        args=training_args,67        train_dataset=tokenized_train,68        eval_dataset=tokenized_eval,69        compute_metrics=compute_metrics,70    )71 72    # 5. Execute Training Loop73    print("🔥 Commencing Neural Network Fine-Tuning...")74    trainer.train()75 76    # 6. Save the Final Model to Disk77    output_path = "../my_fine_tuned_moderator"78    print(f"💾 Saving custom trained model to {output_path}...")79    model.save_pretrained(output_path)80    tokenizer.save_pretrained(output_path)81    print("✅ Training Complete!")82 83if __name__ == "__main__":84    # NOTE FOR STUDENT: Remove 'sample_size=100' when you want to train on the FULL Big Data dataset overnight.85    train_content_moderation(sample_size=100)86