training-scripts / train_tool_calling.py

Upload train_tool_calling.py with huggingface_hub

38d7272 verified 10 months ago

21.7 kB

# /// script
# requires-python = ">=3.10"
# dependencies = [
#     "torch>=2.0.0",
#     "transformers>=4.50.0",
#     "datasets>=2.14.0",
#     "trl>=0.12.0",
#     "peft>=0.7.0",
#     "accelerate>=0.25.0",
#     "bitsandbytes>=0.41.0",
#     "trackio",
#     "huggingface_hub",
# ]
# ///
"""
LoRA Fine-tuning Script: Add Tool Calling to Synthia-S1-27b
This script fine-tunes Tesslate/Synthia-S1-27b with LoRA using the
nvidia/Nemotron-Agentic-v1 tool_calling dataset.
Usage:
    # With uv (recommended)
    uv run train_tool_calling.py
    # Or with pip
    pip install torch transformers datasets trl peft accelerate bitsandbytes trackio
    python train_tool_calling.py
Hardware Requirements:
    - Minimum: 1x A100 80GB or 2x A10G 24GB
    - Recommended: 1x A100 80GB for fastest training
"""

import os
import json
from datasets import load_dataset, Dataset
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, DataCollatorForLanguageModeling
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
from trl import SFTTrainer, SFTConfig
import torch
import trackio
from huggingface_hub import hf_hub_download, HfApi, create_repo

# ============================================================================
# CONFIGURATION - Modify these values as needed
# ============================================================================

# Model configuration
BASE_MODEL = "Tesslate/Synthia-S1-27b"
OUTPUT_MODEL = "Synthia-S1-27b-tool-calling"  # Will be pushed as Codyfederer/Synthia-S1-27b-tool-calling

# Dataset configuration
DATASET_NAME = "nvidia/Nemotron-Agentic-v1"
DATASET_SPLIT = "tool_calling"
MAX_SAMPLES = None  # Set to a number (e.g., 10000) to limit dataset size for testing

# Training hyperparameters
NUM_EPOCHS = 1  # 1 epoch is often sufficient for large datasets
MAX_SEQ_LENGTH = 4096  # Adjust based on your GPU memory
BATCH_SIZE = 1  # Per device batch size
GRADIENT_ACCUMULATION = 16  # Effective batch size = BATCH_SIZE * GRADIENT_ACCUMULATION
LEARNING_RATE = 2e-4
WARMUP_RATIO = 0.03

# LoRA configuration
LORA_R = 64  # LoRA rank - higher = more capacity but more memory
LORA_ALPHA = 128  # LoRA alpha - typically 2x rank
LORA_DROPOUT = 0.05

# Quantization (4-bit for memory efficiency)
USE_4BIT = False  # Using BF16 LoRA on H100 for better quality

# Tokenized dataset caching
TOKENIZED_DATASET_REPO = "Codyfederer/synthia-tool-calling-tokenized"
SAVE_TOKENIZED = True  # Save tokenized dataset to Hub for reuse
TOKENIZED_DATASET_PRIVATE = True  # Make tokenized dataset private
LOAD_TOKENIZED_IF_EXISTS = True  # Skip tokenization if already exists on Hub

# Hub configuration
PUSH_TO_HUB = True
HUB_PRIVATE = False  # Set to True for private model

# ============================================================================
# TRAINING SCRIPT
# ============================================================================

def tokenize_conversation(example, tokenizer, max_length):
    """
    Tokenize a conversation using the model's chat template.
    Returns input_ids, attention_mask, and labels for causal LM training.
    """
    messages = example["messages"]

    # Apply chat template to get the full text
    text = tokenizer.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=False
    )

    # Tokenize the text
    tokenized = tokenizer(
        text,
        truncation=True,
        max_length=max_length,
        padding=False,  # We'll pad later in the data collator
        return_tensors=None,  # Return lists, not tensors
    )

    # For causal LM, labels are the same as input_ids (shifted internally by the model)
    tokenized["labels"] = tokenized["input_ids"].copy()

    return tokenized


def main():
    print("=" * 60)
    print("Tool Calling Fine-tuning for Synthia-S1-27b")
    print("=" * 60)

    # Initialize Trackio for monitoring
    trackio.init(project="synthia-tool-calling")

    # Get HF username for hub_model_id
    from huggingface_hub import whoami
    try:
        username = whoami()["name"]
        hub_model_id = f"{username}/{OUTPUT_MODEL}"
        print(f"Will push to: {hub_model_id}")
    except Exception as e:
        print(f"Warning: Not logged in to HF Hub ({e})")
        print("Model will be saved locally only. Run 'huggingface-cli login' to enable Hub push.")
        hub_model_id = OUTPUT_MODEL
        global PUSH_TO_HUB
        PUSH_TO_HUB = False

    # -------------------------------------------------------------------------
    # Load Tokenizer FIRST (needed for tokenization)
    # -------------------------------------------------------------------------
    print(f"\nLoading tokenizer from {BASE_MODEL}...")

    tokenizer = AutoTokenizer.from_pretrained(
        BASE_MODEL,
        trust_remote_code=True,
        padding_side="right",
    )

    # Ensure pad token is set
    if tokenizer.pad_token is None:
        tokenizer.pad_token = tokenizer.eos_token
        tokenizer.pad_token_id = tokenizer.eos_token_id

    print(f"Vocab size: {len(tokenizer):,}")

    # -------------------------------------------------------------------------
    # Try to Load Pre-tokenized Dataset from Hub
    # -------------------------------------------------------------------------
    train_dataset = None
    eval_dataset = None

    if LOAD_TOKENIZED_IF_EXISTS:
        print(f"\nChecking for pre-tokenized dataset: {TOKENIZED_DATASET_REPO}")
        try:
            from datasets import load_dataset as hf_load_dataset

            # Try to load the tokenized dataset
            tokenized_ds = hf_load_dataset(TOKENIZED_DATASET_REPO)

            # Check if it has the required columns (input_ids, attention_mask)
            if "train" in tokenized_ds and "input_ids" in tokenized_ds["train"].column_names:
                print("  Found pre-tokenized dataset with input_ids!")
                train_dataset = tokenized_ds["train"]
                eval_dataset = tokenized_ds.get("test", tokenized_ds.get("validation"))
                print(f"  Train samples: {len(train_dataset):,}")
                if eval_dataset:
                    print(f"  Eval samples: {len(eval_dataset):,}")
            else:
                print("  Dataset exists but is not tokenized (no input_ids column)")
                print("  Will re-tokenize and save...")
        except Exception as e:
            print(f"  Could not load pre-tokenized dataset: {e}")
            print("  Will tokenize from scratch...")

    # -------------------------------------------------------------------------
    # Load and Tokenize Dataset (if not loaded from Hub)
    # -------------------------------------------------------------------------
    if train_dataset is None:
        print(f"\nLoading dataset: {DATASET_NAME} ({DATASET_SPLIT} split)...")

        # Download the JSONL file directly from the dataset repo
        jsonl_file = f"data/{DATASET_SPLIT}.jsonl"
        print(f"Downloading {jsonl_file}...")

        local_path = hf_hub_download(
            repo_id=DATASET_NAME,
            filename=jsonl_file,
            repo_type="dataset"
        )
        print(f"Downloaded to: {local_path}")

        # Load JSONL manually to handle schema inconsistencies
        print("Loading and processing JSONL file...")
        processed_examples = []
        skipped = 0

        with open(local_path, 'r', encoding='utf-8') as f:
            for line_num, line in enumerate(f):
                if line_num % 50000 == 0:
                    print(f"  Processed {line_num:,} lines...")
                try:
                    example = json.loads(line.strip())
                    messages = example.get("messages", [])

                    # Convert messages to consistent format
                    formatted_messages = []
                    for msg in messages:
                        role = msg.get("role", "user")
                        content = msg.get("content", "")

                        # Handle content that might be a list or complex object
                        if isinstance(content, list):
                            # For tool calls, content is often a list of dicts
                            parts = []
                            for item in content:
                                if isinstance(item, dict):
                                    if "text" in item:
                                        parts.append(item["text"])
                                    else:
                                        parts.append(json.dumps(item))
                                else:
                                    parts.append(str(item))
                            content = "\n".join(parts) if parts else ""
                        elif isinstance(content, dict):
                            content = json.dumps(content)
                        elif content is None:
                            content = ""
                        else:
                            content = str(content)

                        formatted_messages.append({
                            "role": role,
                            "content": content
                        })

                    # Ensure proper role alternation for chat template
                    # Merge consecutive messages with same role, handle tool messages
                    if formatted_messages:
                        merged_messages = []
                        for msg in formatted_messages:
                            role = msg["role"]
                            content = msg["content"]

                            # Map tool role to assistant (tool responses are from assistant's perspective)
                            if role == "tool":
                                role = "user"  # Tool output is provided to the model like user input
                                content = f"[Tool Result]\n{content}"

                            # If same role as previous, merge content
                            if merged_messages and merged_messages[-1]["role"] == role:
                                merged_messages[-1]["content"] += f"\n\n{content}"
                            else:
                                merged_messages.append({"role": role, "content": content})

                        # Ensure conversation starts with user and alternates
                        if merged_messages and merged_messages[0]["role"] != "user":
                            # Prepend a placeholder user message if starts with assistant
                            merged_messages.insert(0, {"role": "user", "content": "[Start]"})

                        processed_examples.append({"messages": merged_messages})

                except Exception as e:
                    skipped += 1
                    if skipped < 5:
                        print(f"  Warning: Skipped line {line_num}: {e}")

        print(f"Loaded {len(processed_examples):,} examples (skipped {skipped})")

        # Create dataset from processed examples
        dataset = Dataset.from_list(processed_examples)
        print(f"Dataset size: {len(dataset):,} examples")

        if MAX_SAMPLES and len(dataset) > MAX_SAMPLES:
            dataset = dataset.shuffle(seed=42).select(range(MAX_SAMPLES))
            print(f"Limited to {MAX_SAMPLES:,} samples for training")

        # Create train/eval split
        split_dataset = dataset.train_test_split(test_size=0.02, seed=42)
        train_dataset = split_dataset["train"]
        eval_dataset = split_dataset["test"]

        print(f"Train samples: {len(train_dataset):,}")
        print(f"Eval samples: {len(eval_dataset):,}")

        # -------------------------------------------------------------------------
        # TOKENIZE the dataset (this is the key step!)
        # -------------------------------------------------------------------------
        print(f"\nTokenizing dataset with max_length={MAX_SEQ_LENGTH}...")
        print("This may take a while for large datasets...")

        # Tokenize train dataset
        train_dataset = train_dataset.map(
            lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
            remove_columns=["messages"],  # Remove text, keep only tokens
            num_proc=4,  # Parallelize
            desc="Tokenizing train",
        )

        # Tokenize eval dataset
        eval_dataset = eval_dataset.map(
            lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
            remove_columns=["messages"],
            num_proc=4,
            desc="Tokenizing eval",
        )

        print(f"Tokenization complete!")
        print(f"Train dataset columns: {train_dataset.column_names}")
        print(f"Sample input_ids length: {len(train_dataset[0]['input_ids'])}")

        # Save TOKENIZED dataset to Hub for reuse
        if SAVE_TOKENIZED:
            print(f"\nSaving TOKENIZED dataset to Hub: {TOKENIZED_DATASET_REPO}")
            try:
                # Create the repo if it doesn't exist (private!)
                api = HfApi()
                try:
                    create_repo(
                        TOKENIZED_DATASET_REPO,
                        repo_type="dataset",
                        private=TOKENIZED_DATASET_PRIVATE,
                        exist_ok=True
                    )
                    print(f"  Created/verified repo (private={TOKENIZED_DATASET_PRIVATE})")

                    # Try to update visibility if repo already exists
                    if TOKENIZED_DATASET_PRIVATE:
                        try:
                            api.update_repo_visibility(
                                TOKENIZED_DATASET_REPO,
                                repo_type="dataset",
                                private=True
                            )
                            print(f"  Ensured repo is private")
                        except Exception:
                            pass  # Ignore if already private or no permission
                except Exception as e:
                    print(f"  Repo creation note: {e}")

                # Reset format to ensure data is serializable (not torch tensors)
                train_dataset.reset_format()
                eval_dataset.reset_format()

                # Verify the data looks correct before pushing
                print(f"  Verifying tokenized data...")
                print(f"    Train columns: {train_dataset.column_names}")
                print(f"    Sample input_ids type: {type(train_dataset[0]['input_ids'])}")
                print(f"    Sample input_ids length: {len(train_dataset[0]['input_ids'])}")
                print(f"    First 10 tokens: {train_dataset[0]['input_ids'][:10]}")

                # Push tokenized datasets to Hub (private is set at repo creation)
                print(f"  Pushing train split ({len(train_dataset):,} examples)...")
                train_dataset.push_to_hub(
                    TOKENIZED_DATASET_REPO,
                    split="train",
                )
                print(f"  Pushing test split ({len(eval_dataset):,} examples)...")
                eval_dataset.push_to_hub(
                    TOKENIZED_DATASET_REPO,
                    split="test",
                )
                print(f"  SUCCESS! Saved TOKENIZED data to: https://huggingface.co/datasets/{TOKENIZED_DATASET_REPO}")
                print(f"  Columns saved: {train_dataset.column_names}")
                print(f"  Dataset is private: {TOKENIZED_DATASET_PRIVATE}")

                # Verify the upload by trying to load it back
                print(f"  Verifying upload...")
                try:
                    from datasets import load_dataset as verify_load
                    verify_ds = verify_load(TOKENIZED_DATASET_REPO, split="train", streaming=True)
                    sample = next(iter(verify_ds))
                    if "input_ids" in sample:
                        print(f"  VERIFIED: Dataset contains input_ids with {len(sample['input_ids'])} tokens")
                    else:
                        print(f"  WARNING: Dataset uploaded but input_ids not found in columns: {list(sample.keys())}")
                except Exception as ve:
                    print(f"  Could not verify upload: {ve}")

            except Exception as e:
                print(f"  ERROR saving to Hub: {e}")
                import traceback
                traceback.print_exc()
                print("  Continuing with training anyway...")

    # -------------------------------------------------------------------------
    # Load Model with Quantization
    # -------------------------------------------------------------------------
    print(f"\nLoading model: {BASE_MODEL}...")

    if USE_4BIT:
        print("Using 4-bit quantization (QLoRA)")
        bnb_config = BitsAndBytesConfig(
            load_in_4bit=True,
            bnb_4bit_quant_type="nf4",
            bnb_4bit_compute_dtype=torch.bfloat16,
            bnb_4bit_use_double_quant=True,
        )
    else:
        bnb_config = None

    model = AutoModelForCausalLM.from_pretrained(
        BASE_MODEL,
        quantization_config=bnb_config,
        device_map="auto",
        trust_remote_code=True,
        torch_dtype=torch.bfloat16,
        attn_implementation="sdpa",  # Use PyTorch's Scaled Dot Product Attention
    )

    if USE_4BIT:
        model = prepare_model_for_kbit_training(model)

    print(f"Model loaded. Parameters: {model.num_parameters():,}")

    # -------------------------------------------------------------------------
    # Configure LoRA
    # -------------------------------------------------------------------------
    print(f"\nConfiguring LoRA (r={LORA_R}, alpha={LORA_ALPHA})...")

    # Target modules for Gemma 3 architecture
    target_modules = [
        "q_proj", "k_proj", "v_proj", "o_proj",  # Attention
        "gate_proj", "up_proj", "down_proj",      # MLP
    ]

    lora_config = LoraConfig(
        r=LORA_R,
        lora_alpha=LORA_ALPHA,
        lora_dropout=LORA_DROPOUT,
        target_modules=target_modules,
        bias="none",
        task_type="CAUSAL_LM",
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # -------------------------------------------------------------------------
    # Training Configuration
    # -------------------------------------------------------------------------
    print("\nConfiguring training...")

    training_args = SFTConfig(
        output_dir=f"./{OUTPUT_MODEL}",

        # Training params
        num_train_epochs=NUM_EPOCHS,
        per_device_train_batch_size=BATCH_SIZE,
        per_device_eval_batch_size=BATCH_SIZE,
        gradient_accumulation_steps=GRADIENT_ACCUMULATION,

        # Optimizer
        learning_rate=LEARNING_RATE,
        lr_scheduler_type="cosine",
        warmup_ratio=WARMUP_RATIO,
        weight_decay=0.01,
        optim="adamw_torch",

        # Memory optimization
        gradient_checkpointing=True,
        gradient_checkpointing_kwargs={"use_reentrant": False},
        max_grad_norm=1.0,

        # Sequence length
        max_length=MAX_SEQ_LENGTH,
        packing=False,  # Disable packing for tool calling (preserve conversation structure)

        # Evaluation
        eval_strategy="steps",
        eval_steps=500,

        # Saving
        save_strategy="steps",
        save_steps=500,
        save_total_limit=3,

        # Hub
        push_to_hub=PUSH_TO_HUB,
        hub_model_id=hub_model_id if PUSH_TO_HUB else None,
        hub_strategy="checkpoint",
        hub_private_repo=HUB_PRIVATE,

        # Logging
        logging_steps=10,
        report_to="trackio",
        run_name=f"lora-r{LORA_R}-lr{LEARNING_RATE}",

        # Performance
        bf16=True,
        dataloader_num_workers=4,
        dataloader_pin_memory=True,

        # Reproducibility
        seed=42,
    )

    # -------------------------------------------------------------------------
    # Initialize Trainer
    # -------------------------------------------------------------------------
    print("\nInitializing trainer...")

    # Create data collator for padding pre-tokenized data
    data_collator = DataCollatorForLanguageModeling(
        tokenizer=tokenizer,
        mlm=False,  # Causal LM, not masked LM
    )

    # Check if dataset is pre-tokenized
    is_pretokenized = "input_ids" in train_dataset.column_names
    print(f"Dataset is pre-tokenized: {is_pretokenized}")
    print(f"Dataset columns: {train_dataset.column_names}")

    trainer = SFTTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        processing_class=tokenizer,
        data_collator=data_collator,
    )

    # -------------------------------------------------------------------------
    # Train!
    # -------------------------------------------------------------------------
    print("\n" + "=" * 60)
    print("Starting training...")
    print("=" * 60 + "\n")

    trainer.train()

    # -------------------------------------------------------------------------
    # Save Final Model
    # -------------------------------------------------------------------------
    print("\nSaving final model...")
    trainer.save_model()

    if PUSH_TO_HUB:
        print(f"Pushing to Hub: {hub_model_id}")
        trainer.push_to_hub()
        print(f"\n✅ Model available at: https://huggingface.co/{hub_model_id}")
    else:
        print(f"Model saved locally to: ./{OUTPUT_MODEL}")

    print("\n" + "=" * 60)
    print("Training complete!")
    print("=" * 60)


if __name__ == "__main__":
    main()

	# /// script
	# requires-python = ">=3.10"
	# dependencies = [
	# "torch>=2.0.0",
	# "transformers>=4.50.0",
	# "datasets>=2.14.0",
	# "trl>=0.12.0",
	# "peft>=0.7.0",
	# "accelerate>=0.25.0",
	# "bitsandbytes>=0.41.0",
	# "trackio",
	# "huggingface_hub",
	# ]
	# ///
	"""
	LoRA Fine-tuning Script: Add Tool Calling to Synthia-S1-27b

	This script fine-tunes Tesslate/Synthia-S1-27b with LoRA using the
	nvidia/Nemotron-Agentic-v1 tool_calling dataset.

	Usage:
	# With uv (recommended)
	uv run train_tool_calling.py

	# Or with pip
	pip install torch transformers datasets trl peft accelerate bitsandbytes trackio
	python train_tool_calling.py

	Hardware Requirements:
	- Minimum: 1x A100 80GB or 2x A10G 24GB
	- Recommended: 1x A100 80GB for fastest training
	"""

	import os
	import json
	from datasets import load_dataset, Dataset
	from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, DataCollatorForLanguageModeling
	from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
	from trl import SFTTrainer, SFTConfig
	import torch
	import trackio
	from huggingface_hub import hf_hub_download, HfApi, create_repo

	# ============================================================================
	# CONFIGURATION - Modify these values as needed
	# ============================================================================

	# Model configuration
	BASE_MODEL = "Tesslate/Synthia-S1-27b"
	OUTPUT_MODEL = "Synthia-S1-27b-tool-calling" # Will be pushed as Codyfederer/Synthia-S1-27b-tool-calling

	# Dataset configuration
	DATASET_NAME = "nvidia/Nemotron-Agentic-v1"
	DATASET_SPLIT = "tool_calling"
	MAX_SAMPLES = None # Set to a number (e.g., 10000) to limit dataset size for testing

	# Training hyperparameters
	NUM_EPOCHS = 1 # 1 epoch is often sufficient for large datasets
	MAX_SEQ_LENGTH = 4096 # Adjust based on your GPU memory
	BATCH_SIZE = 1 # Per device batch size
	GRADIENT_ACCUMULATION = 16 # Effective batch size = BATCH_SIZE * GRADIENT_ACCUMULATION
	LEARNING_RATE = 2e-4
	WARMUP_RATIO = 0.03

	# LoRA configuration
	LORA_R = 64 # LoRA rank - higher = more capacity but more memory
	LORA_ALPHA = 128 # LoRA alpha - typically 2x rank
	LORA_DROPOUT = 0.05

	# Quantization (4-bit for memory efficiency)
	USE_4BIT = False # Using BF16 LoRA on H100 for better quality

	# Tokenized dataset caching
	TOKENIZED_DATASET_REPO = "Codyfederer/synthia-tool-calling-tokenized"
	SAVE_TOKENIZED = True # Save tokenized dataset to Hub for reuse
	TOKENIZED_DATASET_PRIVATE = True # Make tokenized dataset private
	LOAD_TOKENIZED_IF_EXISTS = True # Skip tokenization if already exists on Hub

	# Hub configuration
	PUSH_TO_HUB = True
	HUB_PRIVATE = False # Set to True for private model

	# ============================================================================
	# TRAINING SCRIPT
	# ============================================================================

	def tokenize_conversation(example, tokenizer, max_length):
	"""
	Tokenize a conversation using the model's chat template.
	Returns input_ids, attention_mask, and labels for causal LM training.
	"""
	messages = example["messages"]

	# Apply chat template to get the full text
	text = tokenizer.apply_chat_template(
	messages,
	tokenize=False,
	add_generation_prompt=False
	)

	# Tokenize the text
	tokenized = tokenizer(
	text,
	truncation=True,
	max_length=max_length,
	padding=False, # We'll pad later in the data collator
	return_tensors=None, # Return lists, not tensors
	)

	# For causal LM, labels are the same as input_ids (shifted internally by the model)
	tokenized["labels"] = tokenized["input_ids"].copy()

	return tokenized


	def main():
	print("=" * 60)
	print("Tool Calling Fine-tuning for Synthia-S1-27b")
	print("=" * 60)

	# Initialize Trackio for monitoring
	trackio.init(project="synthia-tool-calling")

	# Get HF username for hub_model_id
	from huggingface_hub import whoami
	try:
	username = whoami()["name"]
	hub_model_id = f"{username}/{OUTPUT_MODEL}"
	print(f"Will push to: {hub_model_id}")
	except Exception as e:
	print(f"Warning: Not logged in to HF Hub ({e})")
	print("Model will be saved locally only. Run 'huggingface-cli login' to enable Hub push.")
	hub_model_id = OUTPUT_MODEL
	global PUSH_TO_HUB
	PUSH_TO_HUB = False

	# -------------------------------------------------------------------------
	# Load Tokenizer FIRST (needed for tokenization)
	# -------------------------------------------------------------------------
	print(f"\nLoading tokenizer from {BASE_MODEL}...")

	tokenizer = AutoTokenizer.from_pretrained(
	BASE_MODEL,
	trust_remote_code=True,
	padding_side="right",
	)

	# Ensure pad token is set
	if tokenizer.pad_token is None:
	tokenizer.pad_token = tokenizer.eos_token
	tokenizer.pad_token_id = tokenizer.eos_token_id

	print(f"Vocab size: {len(tokenizer):,}")

	# -------------------------------------------------------------------------
	# Try to Load Pre-tokenized Dataset from Hub
	# -------------------------------------------------------------------------
	train_dataset = None
	eval_dataset = None

	if LOAD_TOKENIZED_IF_EXISTS:
	print(f"\nChecking for pre-tokenized dataset: {TOKENIZED_DATASET_REPO}")
	try:
	from datasets import load_dataset as hf_load_dataset

	# Try to load the tokenized dataset
	tokenized_ds = hf_load_dataset(TOKENIZED_DATASET_REPO)

	# Check if it has the required columns (input_ids, attention_mask)
	if "train" in tokenized_ds and "input_ids" in tokenized_ds["train"].column_names:
	print(" Found pre-tokenized dataset with input_ids!")
	train_dataset = tokenized_ds["train"]
	eval_dataset = tokenized_ds.get("test", tokenized_ds.get("validation"))
	print(f" Train samples: {len(train_dataset):,}")
	if eval_dataset:
	print(f" Eval samples: {len(eval_dataset):,}")
	else:
	print(" Dataset exists but is not tokenized (no input_ids column)")
	print(" Will re-tokenize and save...")
	except Exception as e:
	print(f" Could not load pre-tokenized dataset: {e}")
	print(" Will tokenize from scratch...")

	# -------------------------------------------------------------------------
	# Load and Tokenize Dataset (if not loaded from Hub)
	# -------------------------------------------------------------------------
	if train_dataset is None:
	print(f"\nLoading dataset: {DATASET_NAME} ({DATASET_SPLIT} split)...")

	# Download the JSONL file directly from the dataset repo
	jsonl_file = f"data/{DATASET_SPLIT}.jsonl"
	print(f"Downloading {jsonl_file}...")

	local_path = hf_hub_download(
	repo_id=DATASET_NAME,
	filename=jsonl_file,
	repo_type="dataset"
	)
	print(f"Downloaded to: {local_path}")

	# Load JSONL manually to handle schema inconsistencies
	print("Loading and processing JSONL file...")
	processed_examples = []
	skipped = 0

	with open(local_path, 'r', encoding='utf-8') as f:
	for line_num, line in enumerate(f):
	if line_num % 50000 == 0:
	print(f" Processed {line_num:,} lines...")
	try:
	example = json.loads(line.strip())
	messages = example.get("messages", [])

	# Convert messages to consistent format
	formatted_messages = []
	for msg in messages:
	role = msg.get("role", "user")
	content = msg.get("content", "")

	# Handle content that might be a list or complex object
	if isinstance(content, list):
	# For tool calls, content is often a list of dicts
	parts = []
	for item in content:
	if isinstance(item, dict):
	if "text" in item:
	parts.append(item["text"])
	else:
	parts.append(json.dumps(item))
	else:
	parts.append(str(item))
	content = "\n".join(parts) if parts else ""
	elif isinstance(content, dict):
	content = json.dumps(content)
	elif content is None:
	content = ""
	else:
	content = str(content)

	formatted_messages.append({
	"role": role,
	"content": content
	})

	# Ensure proper role alternation for chat template
	# Merge consecutive messages with same role, handle tool messages
	if formatted_messages:
	merged_messages = []
	for msg in formatted_messages:
	role = msg["role"]
	content = msg["content"]

	# Map tool role to assistant (tool responses are from assistant's perspective)
	if role == "tool":
	role = "user" # Tool output is provided to the model like user input
	content = f"[Tool Result]\n{content}"

	# If same role as previous, merge content
	if merged_messages and merged_messages[-1]["role"] == role:
	merged_messages[-1]["content"] += f"\n\n{content}"
	else:
	merged_messages.append({"role": role, "content": content})

	# Ensure conversation starts with user and alternates
	if merged_messages and merged_messages[0]["role"] != "user":
	# Prepend a placeholder user message if starts with assistant
	merged_messages.insert(0, {"role": "user", "content": "[Start]"})

	processed_examples.append({"messages": merged_messages})

	except Exception as e:
	skipped += 1
	if skipped < 5:
	print(f" Warning: Skipped line {line_num}: {e}")

	print(f"Loaded {len(processed_examples):,} examples (skipped {skipped})")

	# Create dataset from processed examples
	dataset = Dataset.from_list(processed_examples)
	print(f"Dataset size: {len(dataset):,} examples")

	if MAX_SAMPLES and len(dataset) > MAX_SAMPLES:
	dataset = dataset.shuffle(seed=42).select(range(MAX_SAMPLES))
	print(f"Limited to {MAX_SAMPLES:,} samples for training")

	# Create train/eval split
	split_dataset = dataset.train_test_split(test_size=0.02, seed=42)
	train_dataset = split_dataset["train"]
	eval_dataset = split_dataset["test"]

	print(f"Train samples: {len(train_dataset):,}")
	print(f"Eval samples: {len(eval_dataset):,}")

	# -------------------------------------------------------------------------
	# TOKENIZE the dataset (this is the key step!)
	# -------------------------------------------------------------------------
	print(f"\nTokenizing dataset with max_length={MAX_SEQ_LENGTH}...")
	print("This may take a while for large datasets...")

	# Tokenize train dataset
	train_dataset = train_dataset.map(
	lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
	remove_columns=["messages"], # Remove text, keep only tokens
	num_proc=4, # Parallelize
	desc="Tokenizing train",
	)

	# Tokenize eval dataset
	eval_dataset = eval_dataset.map(
	lambda x: tokenize_conversation(x, tokenizer, MAX_SEQ_LENGTH),
	remove_columns=["messages"],
	num_proc=4,
	desc="Tokenizing eval",
	)

	print(f"Tokenization complete!")
	print(f"Train dataset columns: {train_dataset.column_names}")
	print(f"Sample input_ids length: {len(train_dataset[0]['input_ids'])}")

	# Save TOKENIZED dataset to Hub for reuse
	if SAVE_TOKENIZED:
	print(f"\nSaving TOKENIZED dataset to Hub: {TOKENIZED_DATASET_REPO}")
	try:
	# Create the repo if it doesn't exist (private!)
	api = HfApi()
	try:
	create_repo(
	TOKENIZED_DATASET_REPO,
	repo_type="dataset",
	private=TOKENIZED_DATASET_PRIVATE,
	exist_ok=True
	)
	print(f" Created/verified repo (private={TOKENIZED_DATASET_PRIVATE})")

	# Try to update visibility if repo already exists
	if TOKENIZED_DATASET_PRIVATE:
	try:
	api.update_repo_visibility(
	TOKENIZED_DATASET_REPO,
	repo_type="dataset",
	private=True
	)
	print(f" Ensured repo is private")
	except Exception:
	pass # Ignore if already private or no permission
	except Exception as e:
	print(f" Repo creation note: {e}")

	# Reset format to ensure data is serializable (not torch tensors)
	train_dataset.reset_format()
	eval_dataset.reset_format()

	# Verify the data looks correct before pushing
	print(f" Verifying tokenized data...")
	print(f" Train columns: {train_dataset.column_names}")
	print(f" Sample input_ids type: {type(train_dataset[0]['input_ids'])}")
	print(f" Sample input_ids length: {len(train_dataset[0]['input_ids'])}")
	print(f" First 10 tokens: {train_dataset[0]['input_ids'][:10]}")

	# Push tokenized datasets to Hub (private is set at repo creation)
	print(f" Pushing train split ({len(train_dataset):,} examples)...")
	train_dataset.push_to_hub(
	TOKENIZED_DATASET_REPO,
	split="train",
	)
	print(f" Pushing test split ({len(eval_dataset):,} examples)...")
	eval_dataset.push_to_hub(
	TOKENIZED_DATASET_REPO,
	split="test",
	)
	print(f" SUCCESS! Saved TOKENIZED data to: https://huggingface.co/datasets/{TOKENIZED_DATASET_REPO}")
	print(f" Columns saved: {train_dataset.column_names}")
	print(f" Dataset is private: {TOKENIZED_DATASET_PRIVATE}")

	# Verify the upload by trying to load it back
	print(f" Verifying upload...")
	try:
	from datasets import load_dataset as verify_load
	verify_ds = verify_load(TOKENIZED_DATASET_REPO, split="train", streaming=True)
	sample = next(iter(verify_ds))
	if "input_ids" in sample:
	print(f" VERIFIED: Dataset contains input_ids with {len(sample['input_ids'])} tokens")
	else:
	print(f" WARNING: Dataset uploaded but input_ids not found in columns: {list(sample.keys())}")
	except Exception as ve:
	print(f" Could not verify upload: {ve}")

	except Exception as e:
	print(f" ERROR saving to Hub: {e}")
	import traceback
	traceback.print_exc()
	print(" Continuing with training anyway...")

	# -------------------------------------------------------------------------
	# Load Model with Quantization
	# -------------------------------------------------------------------------
	print(f"\nLoading model: {BASE_MODEL}...")

	if USE_4BIT:
	print("Using 4-bit quantization (QLoRA)")
	bnb_config = BitsAndBytesConfig(
	load_in_4bit=True,
	bnb_4bit_quant_type="nf4",
	bnb_4bit_compute_dtype=torch.bfloat16,
	bnb_4bit_use_double_quant=True,
	)
	else:
	bnb_config = None

	model = AutoModelForCausalLM.from_pretrained(
	BASE_MODEL,
	quantization_config=bnb_config,
	device_map="auto",
	trust_remote_code=True,
	torch_dtype=torch.bfloat16,
	attn_implementation="sdpa", # Use PyTorch's Scaled Dot Product Attention
	)

	if USE_4BIT:
	model = prepare_model_for_kbit_training(model)

	print(f"Model loaded. Parameters: {model.num_parameters():,}")

	# -------------------------------------------------------------------------
	# Configure LoRA
	# -------------------------------------------------------------------------
	print(f"\nConfiguring LoRA (r={LORA_R}, alpha={LORA_ALPHA})...")

	# Target modules for Gemma 3 architecture
	target_modules = [
	"q_proj", "k_proj", "v_proj", "o_proj", # Attention
	"gate_proj", "up_proj", "down_proj", # MLP
	]

	lora_config = LoraConfig(
	r=LORA_R,
	lora_alpha=LORA_ALPHA,
	lora_dropout=LORA_DROPOUT,
	target_modules=target_modules,
	bias="none",
	task_type="CAUSAL_LM",
	)

	model = get_peft_model(model, lora_config)
	model.print_trainable_parameters()

	# -------------------------------------------------------------------------
	# Training Configuration
	# -------------------------------------------------------------------------
	print("\nConfiguring training...")

	training_args = SFTConfig(
	output_dir=f"./{OUTPUT_MODEL}",

	# Training params
	num_train_epochs=NUM_EPOCHS,
	per_device_train_batch_size=BATCH_SIZE,
	per_device_eval_batch_size=BATCH_SIZE,
	gradient_accumulation_steps=GRADIENT_ACCUMULATION,

	# Optimizer
	learning_rate=LEARNING_RATE,
	lr_scheduler_type="cosine",
	warmup_ratio=WARMUP_RATIO,
	weight_decay=0.01,
	optim="adamw_torch",

	# Memory optimization
	gradient_checkpointing=True,
	gradient_checkpointing_kwargs={"use_reentrant": False},
	max_grad_norm=1.0,

	# Sequence length
	max_length=MAX_SEQ_LENGTH,
	packing=False, # Disable packing for tool calling (preserve conversation structure)

	# Evaluation
	eval_strategy="steps",
	eval_steps=500,

	# Saving
	save_strategy="steps",
	save_steps=500,
	save_total_limit=3,

	# Hub
	push_to_hub=PUSH_TO_HUB,
	hub_model_id=hub_model_id if PUSH_TO_HUB else None,
	hub_strategy="checkpoint",
	hub_private_repo=HUB_PRIVATE,

	# Logging
	logging_steps=10,
	report_to="trackio",
	run_name=f"lora-r{LORA_R}-lr{LEARNING_RATE}",

	# Performance
	bf16=True,
	dataloader_num_workers=4,
	dataloader_pin_memory=True,

	# Reproducibility
	seed=42,
	)

	# -------------------------------------------------------------------------
	# Initialize Trainer
	# -------------------------------------------------------------------------
	print("\nInitializing trainer...")

	# Create data collator for padding pre-tokenized data
	data_collator = DataCollatorForLanguageModeling(
	tokenizer=tokenizer,
	mlm=False, # Causal LM, not masked LM
	)

	# Check if dataset is pre-tokenized
	is_pretokenized = "input_ids" in train_dataset.column_names
	print(f"Dataset is pre-tokenized: {is_pretokenized}")
	print(f"Dataset columns: {train_dataset.column_names}")

	trainer = SFTTrainer(
	model=model,
	args=training_args,
	train_dataset=train_dataset,
	eval_dataset=eval_dataset,
	processing_class=tokenizer,
	data_collator=data_collator,
	)

	# -------------------------------------------------------------------------
	# Train!
	# -------------------------------------------------------------------------
	print("\n" + "=" * 60)
	print("Starting training...")
	print("=" * 60 + "\n")

	trainer.train()

	# -------------------------------------------------------------------------
	# Save Final Model
	# -------------------------------------------------------------------------
	print("\nSaving final model...")
	trainer.save_model()

	if PUSH_TO_HUB:
	print(f"Pushing to Hub: {hub_model_id}")
	trainer.push_to_hub()
	print(f"\n✅ Model available at: https://huggingface.co/{hub_model_id}")
	else:
	print(f"Model saved locally to: ./{OUTPUT_MODEL}")

	print("\n" + "=" * 60)
	print("Training complete!")
	print("=" * 60)


	if __name__ == "__main__":
	main()