ysn-rfd/text-dataset-tiny-code-script-py-format

1#!/usr/bin/env python
2import os
3import logging
4import torch
5from transformers import (
6    AutoTokenizer,
7    AutoModelForCausalLM,
8    Trainer,
9    TrainingArguments,
10    DataCollatorForLanguageModeling,
11    get_cosine_schedule_with_warmup,
12)
13from datasets import load_dataset
14
15# Setup logging for progress messages
16logging.basicConfig(level=logging.INFO)
17logger = logging.getLogger(__name__)
18
19# 1. Load the pre-trained tokenizer and model
20model_name = "sshleifer/tiny-gpt2"  # Using GPT-2 as a small language model example
21tokenizer = AutoTokenizer.from_pretrained(model_name)
22
23# Check if a padding token is defined; if not, set it.
24if tokenizer.pad_token is None:
25    # Option 1: Use the end-of-sequence token as the padding token.
26    tokenizer.pad_token = tokenizer.eos_token
27    # Option 2 (uncomment to use a dedicated PAD token):
28    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
29    # After adding special tokens, resize model embeddings:
30    # model.resize_token_embeddings(len(tokenizer))
31
32# Load the pre-trained model.
33model = AutoModelForCausalLM.from_pretrained(model_name)
34
35# 2. Prepare the dataset
36# For demonstration, we use the Wikitext-2 raw dataset.
37dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
38
39# Define a tokenization function.
40def tokenize_function(examples):
41    # Tokenize texts with padding (to max_length) and truncation.
42    # Here, we set max_length=32. Adjust as needed.
43    return tokenizer(
44        examples["text"],
45        truncation=True,
46        max_length=32,
47        padding="max_length"
48    )
49
50# Apply the tokenization function over the dataset.
51tokenized_dataset = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
52
53# 3. Create a data collator for language modeling (no masked LM).
54data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
55
56# 4. Setup training arguments
57training_args = TrainingArguments(
58    output_dir="./gpt2-finetuned",
59    overwrite_output_dir=True,
60    num_train_epochs=1,              # Adjust the number of epochs as needed
61    per_device_train_batch_size=8,   # Adjust based on your available GPU memory
62    save_steps=1000,
63    save_total_limit=2,
64    logging_steps=100,
65    prediction_loss_only=True,       # Useful for language modeling tasks
66)
67
68# 5. Create the custom optimizer and scheduler
69# Calculate the total number of training steps
70num_update_steps_per_epoch = len(tokenized_dataset) // training_args.per_device_train_batch_size
71max_train_steps = training_args.num_train_epochs * num_update_steps_per_epoch
72
73# Create AdamW optimizer with a custom learning rate and weight decay.
74optimizer = torch.optim.AdamW(model.parameters(), lr=0.1, weight_decay=0.1)
75
76# Create a cosine learning rate scheduler with warmup.
77scheduler = get_cosine_schedule_with_warmup(
78    optimizer,
79    num_warmup_steps=100,       # Number of warmup steps, adjust as needed
80    num_training_steps=max_train_steps
81)
82
83# 6. Initialize the Trainer with the custom optimizer and scheduler.
84trainer = Trainer(
85    model=model,
86    args=training_args,
87    train_dataset=tokenized_dataset,
88    data_collator=data_collator,
89    optimizers=(optimizer, scheduler)  # Pass the optimizer and scheduler as a tuple.
90)
91
92# 7. Start training
93logger.info("Starting training...")
94trainer.train()
95
96# 8. Save the fine-tuned model and tokenizer
97model.save_pretrained("./gpt2-finetuned")
98tokenizer.save_pretrained("./gpt2-finetuned")
99logger.info("Training complete and model saved.")
100