Making MCP Tool Calls Smarter: A Practical Fine-Tuning Blueprint

from datasets import load_dataset

DATASET_REPO_ID = "yashsoni78/conversation_data_mcp_100"

try:
    # This single command downloads and loads the dataset into memory
    dataset = load_dataset(DATASET_REPO_ID, split="train")
    logging.info(f"✅ Successfully loaded dataset from Hugging Face Hub: {DATASET_REPO_ID}")
except Exception as e:
    logging.error(f"❌ Failed to load dataset from Hugging Face. Please check the repository ID and your connection. Error: {e}")
    exit()

def format_for_chat_template(example):
    """
    Converts the 'conversations' list into the 'messages' format
    that the SFTTrainer expects for chat templates.
    """
    messages = []
    # The first turn is always the system prompt
    if example['conversations'][0]['from'] == 'system':
        messages.append({"role": "system", "content": example['conversations'][0]['value']})
        turns = example['conversations'][1:]
    else:
        turns = example['conversations']

    # Process the rest of the turns
    for turn in turns:
        role = "user" if turn['from'] == 'human' else "assistant"
        messages.append({"role": role, "content": turn['value']})
        
    return {"messages": messages}

# Use the .map() method for efficient processing
chat_dataset = dataset.map(format_for_chat_template)
# Remove the old columns to keep the dataset clean
chat_dataset = chat_dataset.remove_columns(dataset.column_names)

logging.info("✅ Dataset formatted successfully for the trainer.")
print("\n--- Example of one entry in the new format ---")
print(chat_dataset[0]['messages'])

# Model and tokenizer names
model_name = "mistralai/Mistral-7B-Instruct-v0.2"

# Quantization configuration
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16,
    bnb_4bit_use_double_quant=False,
)
logging.info("BitsAndBytesConfig created.")

# Load model
try:
    model = AutoModelForCausalLM.from_pretrained(
        model_name,
        quantization_config=bnb_config,
        torch_dtype=torch.bfloat16,
        device_map="auto",
        trust_remote_code=True,
    )
    model.config.use_cache = False
    model.config.pretraining_tp = 1
    logging.info("Model loaded successfully.")
except Exception as e:
    logging.error(f"Error loading model: {e}")
    exit()

# Load tokenizer
try:
    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "right"
    logging.info("Tokenizer loaded successfully.")
except Exception as e:
    logging.error(f"Error loading tokenizer: {e}")
    exit()

# PEFT configuration
peft_config = LoraConfig(
    lora_alpha=16,
    lora_dropout=0.1,
    r=64,
    bias="none",
    task_type="CAUSAL_LM",
)
logging.info("LoRA config created.")

# Prepare model for training
model = prepare_model_for_kbit_training(model)
model = get_peft_model(model, peft_config)
logging.info("Model prepared for k-bit training with PEFT.")

# Define the directory for TensorBoard logs
tensorboard_log_dir = "./logs"

# Training arguments
training_arguments = TrainingArguments(
    output_dir="./results",          # Directory to save model checkpoints and logs
    num_train_epochs=10,             # Total number of training epochs
    per_device_train_batch_size=2,   # Batch size per device (GPU/CPU) during training
    gradient_accumulation_steps=1,   # Number of steps to accumulate gradients before updating weights
    optim="paged_adamw_32bit",       # Optimizer type; here, a memory-efficient 32-bit AdamW variant
    save_steps=50,                   # Save checkpoint every 50 steps
    logging_steps=10,                # Log training metrics every 10 steps
    logging_strategy="steps",        # Logging strategy: log every few steps
    logging_dir=tensorboard_log_dir, # Directory for TensorBoard logs
    learning_rate=2e-4,              # Learning rate for optimizer
    weight_decay=0.001,              # Weight decay (L2 regularization) to avoid overfitting
    fp16=False,                      # Use 16-bit floating point precision (disabled)
    bf16=True,                       # Use bfloat16 precision (enabled), useful on hardware that supports it
    max_grad_norm=0.3,               # Maximum gradient norm for gradient clipping (prevents exploding gradients)
    max_steps=-1,                    # Total number of training steps (-1 means use all steps from `num_train_epochs`)
    warmup_ratio=0.03,               # Fraction of total steps used for learning rate warmup
    group_by_length=True,            # Group sequences of similar lengths to speed up training and reduce padding
    lr_scheduler_type="constant",    # Learning rate scheduler type (constant: no decay)
)

# Initialize the SFTTrainer
trainer = SFTTrainer(
    model=model,                # The model to be fine-tuned
    train_dataset=chat_dataset, # Dataset used for training
    peft_config=peft_config,    # Parameter-efficient fine-tuning (PEFT) configuration
    args=training_arguments,    # Training arguments defined above
)

# Start fine tuning
trainer.train()