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

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