pathii/css_design_snippets

1from datasets import load_dataset2from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments3 4# Load dataset from Hugging Face Hub5dataset = load_dataset("pathii/css_design_snippets")6 7# Load pre-trained model and tokenizer8model_name = "TinyLlama/TinyLlama_v1.1"9model = AutoModelForCausalLM.from_pretrained(model_name)10tokenizer = AutoTokenizer.from_pretrained(model_name)11 12# Tokenize dataset13def tokenize_function(example):14    return tokenizer(example["input"], truncation=True)15 16tokenized_datasets = dataset.map(tokenize_function, batched=True)17 18# Define training arguments19training_args = TrainingArguments(20    output_dir="./model",21    evaluation_strategy="epoch",22    learning_rate=2e-5,23    per_device_train_batch_size=8,24    per_device_eval_batch_size=8,25    num_train_epochs=3,26    weight_decay=0.01,27    save_total_limit=2,28    save_strategy="epoch"29)30 31# Create Trainer32trainer = Trainer(33    model=model,34    args=training_args,35    train_dataset=tokenized_datasets["train"],36    eval_dataset=tokenized_datasets["validation"],37)38 39# Start training40trainer.train()41