!pip install transformers datasets accelerate

from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
from datasets import load_dataset

# Load tiny subset of an instruction dataset
dataset = load_dataset("yahma/alpaca-cleaned", split="train[:100]")  # 100 examples only

# Initialize tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("facebook/MobileLLM-R1-950M")
model = AutoModelForCausalLM.from_pretrained("facebook/MobileLLM-R1-950M")

# Tokenize dataset
def tokenize_function(examples):
    return tokenizer(examples['instruction'] + " " + examples['output'], 
                     truncation=True, max_length=128)

tokenized_dataset = dataset.map(tokenize_function, batched=True)

# Use simple data collator
data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)

# Training arguments (very short for fast training)
training_args = TrainingArguments(
    output_dir="./tiny-model",
    per_device_train_batch_size=4,
    num_train_epochs=1,
    logging_steps=10,
    save_steps=50,
    save_total_limit=1,
    fp16=True,
    remove_unused_columns=False,
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

# Train (should take a few seconds on Colab GPU)
trainer.train()