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

1#!/usr/bin/env python
2import os
3import logging
4from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, DataCollatorForLanguageModeling
5from datasets import load_dataset
6
7# Set up logging
8logging.basicConfig(level=logging.INFO)
9logger = logging.getLogger(__name__)
10
11# 1. Load the pre-trained tokenizer and model
12# Here we use "gpt2" which is a small, open‐source language model.
13model_name = "sshleifer/tiny-gpt2"
14tokenizer = AutoTokenizer.from_pretrained(model_name)
15# For causal language modeling (i.e. text generation)
16model = AutoModelForCausalLM.from_pretrained(model_name)
17
18# 2. Prepare the dataset
19# For demonstration, we use a subset of the WikiText dataset.
20# You can replace this with your own text data.
21dataset = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
22
23# Tokenize the texts
24def tokenize_function(examples):
25    return tokenizer(examples["text"], truncation=True, max_length=256)
26
27tokenized_datasets = dataset.map(tokenize_function, batched=True, remove_columns=["text"])
28
29# Create a data collator that handles dynamic padding and masks for language modeling.
30data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
31
32# 3. Set up training arguments
33training_args = TrainingArguments(
34    output_dir="./gpt2-finetuned",
35    overwrite_output_dir=True,
36    num_train_epochs=3,              # Adjust as needed
37    per_device_train_batch_size=4,   # Adjust based on your GPU memory
38    save_steps=500,
39    save_total_limit=2,
40    logging_steps=100,
41    prediction_loss_only=True,       # Only compute loss for LM tasks
42)
43
44# 4. Initialize the Trainer
45trainer = Trainer(
46    model=model,
47    args=training_args,
48    train_dataset=tokenized_datasets,
49    data_collator=data_collator,
50)
51
52# 5. Start training
53logger.info("Starting training...")
54trainer.train()
55
56# 6. Save the model and tokenizer
57model.save_pretrained("./gpt2-finetuned")
58tokenizer.save_pretrained("./gpt2-finetuned")
59logger.info("Training complete and model saved.")
60