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_intent_classifier(sample_size=None):19    print("🚀 Starting Intent Classification Model Training Pipeline...")20    21    # 1. Load a Massive Customer Support Dataset22    # Bitext is a famous customer service intent dataset with thousands of rows23    print("⬇️ Downloading Customer Support Dataset...")24    dataset = load_dataset("bitext/Bitext-customer-support-llm-chatbot-training-dataset")25    26    # The bitext dataset has 'instruction' and 'intent'. We must map string intents to ID integers.27    unique_intents = list(set(dataset['train']['intent']))28    intent2id = {intent: idx for idx, intent in enumerate(unique_intents)}29    id2intent = {idx: intent for intent, idx in intent2id.items()}30 31    def map_labels(example):32        example['label'] = intent2id[example['intent']]33        return example34 35    dataset = dataset.map(map_labels)36 37    # For training validation split (since bitext only has 'train' split initially)38    dataset = dataset['train'].train_test_split(test_size=0.1)39 40    if sample_size:41        print(f"⚠️ Demo Mode: Only training on {sample_size} samples for speed.")42        train_dataset = dataset["train"].select(range(sample_size))43        eval_dataset = dataset["test"].select(range(min(sample_size, len(dataset["test"]))))44    else:45        print("🧠 Big Data Mode: Training on full dataset!")46        train_dataset = dataset["train"]47        eval_dataset = dataset["test"]48 49    # 2. Tokenization50    model_name = "distilbert-base-uncased"51    print(f"🔨 Initializing Tokenizer: {model_name}")52    tokenizer = AutoTokenizer.from_pretrained(model_name)53 54    def tokenize_function(examples):55        return tokenizer(examples["instruction"], padding="max_length", truncation=True, max_length=128)56 57    tokenized_train = train_dataset.map(tokenize_function, batched=True)58    tokenized_eval = eval_dataset.map(tokenize_function, batched=True)59 60    # 3. Model Initialization61    print(f"🤖 Loading DistilBERT for {len(unique_intents)} unique intents...")62    model = AutoModelForSequenceClassification.from_pretrained(63        model_name, 64        num_labels=len(unique_intents),65        id2label=id2intent,66        label2id=intent2id67    )68 69    # 4. Training Configuration70    training_args = TrainingArguments(71        output_dir="./results_intent",72        evaluation_strategy="epoch",73        learning_rate=3e-5,74        per_device_train_batch_size=32,75        per_device_eval_batch_size=32,76        num_train_epochs=4 if not sample_size else 1,77        weight_decay=0.01,78        save_strategy="epoch",79        load_best_model_at_end=True,80    )81 82    trainer = Trainer(83        model=model,84        args=training_args,85        train_dataset=tokenized_train,86        eval_dataset=tokenized_eval,87        compute_metrics=compute_metrics,88    )89 90    # 5. Execute Training Loop91    print("🔥 Commencing Neural Network Fine-Tuning for Customer Intents...")92    trainer.train()93 94    # 6. Save the Final Model to Disk95    output_path = "../my_fine_tuned_intent"96    print(f"💾 Saving custom trained model to {output_path}...")97    model.save_pretrained(output_path)98    tokenizer.save_pretrained(output_path)99    print("✅ Training Complete!")100 101if __name__ == "__main__":102    # NOTE FOR STUDENT: Remove 'sample_size=100' when you want to train on the FULL dataset.103    train_intent_classifier(sample_size=100)104