MacLeanLuke/gemma-2b-tool-tuned-code

1# Ensure Apple Metal (MPS) is enabled2import torch3import os4from transformers import AutoModelForCausalLM, AutoTokenizer, set_seed5from datasets import load_dataset6from peft import LoraConfig, TaskType7from trl import SFTConfig, SFTTrainer8from enum import Enum9 10# ✅ Set device to Metal Performance Shaders (MPS) for Mac M311device = "mps" if torch.backends.mps.is_available() else "cpu"12print(f"Using device: {device}")13 14# ✅ Set seed for reproducibility15set_seed(42)16 17# ✅ Model and dataset18model_name = "google/gemma-2-2b-it"19dataset_name = "Jofthomas/hermes-function-calling-thinking-V1"20tokenizer = AutoTokenizer.from_pretrained(model_name, token=True)21 22# ✅ Adjust tokenizer with special tokens23class ChatmlSpecialTokens(str, Enum):24    tools = "<tools>"25    eotools = "</tools>"26    think = "<think>"27    eothink = "</think>"28    tool_call="<tool_call>"29    eotool_call="</tool_call>"30    tool_response="<tool_response>"31    eotool_response="</tool_response>"32    pad_token = "<pad>"33    eos_token = "<eos>"34 35    @classmethod36    def list(cls):37        return [c.value for c in cls]38 39tokenizer = AutoTokenizer.from_pretrained(40    model_name,41    pad_token=ChatmlSpecialTokens.pad_token.value,42    additional_special_tokens=ChatmlSpecialTokens.list()43)44 45# ✅ Load model and move it to MPS46model = AutoModelForCausalLM.from_pretrained(model_name, token=True, attn_implementation="eager")47model.resize_token_embeddings(len(tokenizer))48model.to(device)49 50# ✅ Data preprocessing function51def preprocess(sample):52    messages = sample["messages"]53 54    if not messages or not isinstance(messages, list):55        return {"text": ""}  # Return empty text if messages are missing56 57    first_message = messages[0]58 59    # Ensure system messages are merged with the first user message60    if first_message["role"] == "system":61        system_message_content = first_message.get("content", "")62        if len(messages) > 1 and messages[1]["role"] == "user":63            messages[1]["content"] = (64                    system_message_content65                    + "\n\nAlso, before making a call to a function, take the time to plan the function to take. "66                    + "Make that thinking process between <think>{your thoughts}</think>\n\n"67                    + messages[1].get("content", "")68            )69            messages.pop(0)  # Remove system message70 71    # Ensure the conversation alternates between "user" and "assistant"72    valid_roles = ["user", "assistant"]73    cleaned_messages = [74        msg for msg in messages if msg.get("role") in valid_roles and msg.get("content")75    ]76 77    # Check if messages are empty after cleanup78    if not cleaned_messages or cleaned_messages[0]["role"] != "user":79        return {"text": ""}  # Ensure the first message is always from the user80 81    # Apply chat template82    try:83        formatted_text = tokenizer.apply_chat_template(cleaned_messages, tokenize=False)84        return {"text": formatted_text}85    except Exception as e:86        print(f"Error processing message: {e}")87        return {"text": ""}88 89# ✅ Load dataset90dataset = load_dataset(dataset_name, cache_dir="/tmp")91dataset = dataset.rename_column("conversations", "messages")92dataset = dataset.map(preprocess, remove_columns=["messages"])93dataset = dataset["train"].train_test_split(0.1)94 95# ✅ Print dataset size before training96print(f"Training dataset size: {len(dataset['train'])} samples")97print(f"Evaluation dataset size: {len(dataset['test'])} samples")98 99# ✅ LoRA configuration100peft_config = LoraConfig(101    r=16,102    lora_alpha=64,103    lora_dropout=0.05,104    target_modules=["gate_proj", "q_proj", "lm_head", "o_proj", "k_proj", "embed_tokens", "down_proj", "up_proj", "v_proj"],105    task_type=TaskType.CAUSAL_LM,106    bias="none",107)108 109# ✅ Training configuration (adjusted for performance on Mac M3 Max)110num_train_epochs = 5  # ✅ Increase to 5 epochs for better training111max_steps = 1000  # ✅ Ensure at least 1000 training steps112learning_rate = 5e-5  # ✅ Reduce learning rate to prevent overfitting113 114training_arguments = SFTConfig(115    output_dir="gemma-2-2B-it-macM3",116    per_device_train_batch_size=2,  # ✅ Keep small if training on MPS117    per_device_eval_batch_size=2,118    gradient_accumulation_steps=4,  # ✅ Helps fit larger batch sizes119    save_strategy="epoch",120    save_total_limit=2,121    save_safetensors=False,122    evaluation_strategy="epoch",123    logging_steps=5,124    learning_rate=learning_rate,125    max_grad_norm=1.0,126    weight_decay=0.1,127    warmup_ratio=0.1,128    lr_scheduler_type="cosine",129    report_to="tensorboard",130    bf16=True,  # ✅ Efficient mixed precision training for Mac MPS131    push_to_hub=False,132    num_train_epochs=num_train_epochs,133    max_steps=max_steps,  # ✅ Ensure training runs for at least 1000 steps134    gradient_checkpointing=True,135    gradient_checkpointing_kwargs={"use_reentrant": False},136    packing=True,137    max_seq_length=1500,138)139 140# ✅ Trainer setup141trainer = SFTTrainer(142    model=model,143    args=training_arguments,144    train_dataset=dataset["train"],145    eval_dataset=dataset["test"],146    processing_class=tokenizer,147    peft_config=peft_config,148)149 150# ✅ Start training (should work efficiently on Mac M3 Max)151trainer.train()152trainer.save_model()153 154print("Training complete! 🚀 Model saved successfully.")