pretrain_bert.py

from transformers import (
    BertConfig,
    BertForMaskedLM,
    BertTokenizerFast,
    DataCollatorForLanguageModeling,
    Trainer,
    TrainingArguments,
)
from datasets import load_dataset


tokenizer = BertTokenizerFast(vocab_file="opcode_vocab.txt")

def tokenize_function(examples):
    if tokenizer.pad_token is None:
        tokenizer.add_special_tokens({'pad_token': '[PAD]'})
    return tokenizer(examples["text"], padding="max_length", truncation=True, max_length=512)

print("Tokenizer loaded.")
print("Tokenizer vocab size:", tokenizer.vocab_size)

dataset = load_dataset(
    'text',
    data_files={'train': 'opcodes_300/*.opcodes.txt'}
)

split_dataset = dataset["train"].train_test_split(test_size=0.1, seed=42)



tokenized_dataset = split_dataset.map(
    tokenize_function,
    batched=True,
    remove_columns=["text"],
    num_proc=24
)
print("Dataset tokenized.")

config = BertConfig(
    vocab_size=tokenizer.vocab_size,
    max_position_embeddings=512,
    num_attention_heads=12,
    num_hidden_layers=6,
    type_vocab_size=1,
)

print("Model config vocab size:", config.vocab_size)


model = BertForMaskedLM(config=config)

# Check if the pad token was added in any process and resize if needed.
# It's better to be explicit here before creating the model.
if tokenizer.pad_token is None:
    tokenizer.add_special_tokens({'pad_token': '[PAD]'})
model.resize_token_embeddings(len(tokenizer))
print(f"Model initialized with {model.num_parameters():,} parameters.")

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer,
    mlm=True,
    mlm_probability=0.15 # Standard 15% masking probability
)

training_args = TrainingArguments(
       output_dir="./opcode_bert",
       overwrite_output_dir=True,
       num_train_epochs=5,
       per_device_train_batch_size=16,
       save_steps=2000,
       save_total_limit=3,
       logging_dir='./logs',
       logging_strategy="steps",
       logging_steps=500,
       eval_strategy="steps",
       eval_steps=500,
       load_best_model_at_end=True,
       metric_for_best_model="eval_loss",
   )

trainer = Trainer(
    model=model,
    args=training_args,
    data_collator=data_collator,
    train_dataset=tokenized_dataset["train"],
    eval_dataset=tokenized_dataset["test"],
)

print("Starting training...")
trainer.train()

# Save the final model
trainer.save_model("./opcode_bert_final")
tokenizer.save_pretrained("./opcode_bert_final")
print("Training complete and final model saved.")