File size: 2,047 Bytes
ea3d866 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 | from transformers import AutoTokenizer, AutoModelForCausalLM, Trainer, TrainingArguments, TextDataset, DataCollatorForLanguageModeling
import torch
# Load the tokenizer and model
model_name = "aubmindlab/aragpt2-base" # Arabic GPT-2 model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForCausalLM.from_pretrained(model_name)
# Add a padding token if not present
if tokenizer.pad_token is None:
tokenizer.add_special_tokens({'pad_token': '[PAD]'}) # Adds a [PAD] token if it's not defined
model.resize_token_embeddings(len(tokenizer)) # Resize model's embeddings to account for the new token
# Load and prepare the data for training
def load_dataset(file_path, tokenizer, block_size=128):
# Read the text data from the file
with open(file_path, "r", encoding="utf-8") as f:
text = f.read()
# Tokenize the text data
inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=block_size)
# Return as a PyTorch dataset
dataset = TextDataset(tokenizer=tokenizer, file_path=file_path, block_size=block_size)
return dataset
# Define the path to your text data file
data_file = "data/raw_data/iraqi_data.txt" # Update this path if necessary
train_dataset = load_dataset(data_file, tokenizer)
# Set up training arguments
training_args = TrainingArguments(
output_dir="./models",
overwrite_output_dir=True,
num_train_epochs=3,
per_device_train_batch_size=2,
save_steps=10_000,
save_total_limit=2,
)
# Define the data collator for language modeling
data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)
# Initialize the Trainer
trainer = Trainer(
model=model,
args=training_args,
train_dataset=train_dataset,
data_collator=data_collator,
)
# Save the trained model and tokenizer
trainer.save_model("models/iraqi_proverb_model")
tokenizer.save_pretrained("models/iraqi_proverb_model")
print("Training complete. Model and tokenizer saved to 'models/iraqi_proverb_model'.")
|