File size: 5,757 Bytes
ba3ecf1 | 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 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 | import os
from sklearn.model_selection import train_test_split
import numpy as np
import matplotlib.pyplot as plt
import logging
import sys
import traceback
from datetime import datetime
from datasets import load_dataset
from sentence_transformers import SentenceTransformer, losses
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator,TripletEvaluator,SequentialEvaluator
from sentence_transformers.similarity_functions import SimilarityFunction
from sentence_transformers.trainer import SentenceTransformerTrainer
from sentence_transformers.training_args import SentenceTransformerTrainingArguments
import torch
from sentence_transformers.training_args import BatchSamplers
import pandas as pd
from arabert.preprocess import ArabertPreprocessor
from pathlib import Path
# Configure logging to write to logs.txt
logging.basicConfig(
format="%(asctime)s - %(message)s",
datefmt="%Y-%m-%d %H:%M:%S",
level=logging.INFO,
handlers=[
logging.FileHandler("logs.txt")
]
)
class Tee:
def __init__(self, *files):
self.files = files
def write(self, obj):
for f in self.files:
f.write(obj)
f.flush()
def flush(self):
for f in self.files:
f.flush()
def isatty(self):
return False
log_file = open("logs.txt", "a")
sys.stdout = Tee(sys.stdout, log_file)
train_batch_size = 64
model_name = "bert-base-arabertv02"
model_nickname = "arabert"
timestamp = datetime.now().strftime("%Y%m%d_%H%M")
output_dir = f"output/{model_nickname}_{timestamp}"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SentenceTransformer(model_name, device=device)
print(f"Model is running on: {device}")
logging.info("Reading the training and eval dataset")
train_dataset = load_dataset("csv", data_files="train.csv")
eval_dataset = load_dataset("csv", data_files="val.csv")
test_dataset = load_dataset("csv", data_files="test.csv")
logging.info(train_dataset)
logging.info(eval_dataset)
logging.info(test_dataset)
# Add this line before the 'evaluators = []' loop
eval_subset = eval_dataset["train"].shuffle(seed=42).select(range(250000))
# Training loss
matryoshka_dims = [768, 512, 256, 128, 64]
inner_train_loss = losses.MultipleNegativesRankingLoss(model=model)
train_loss = losses.MatryoshkaLoss(model, inner_train_loss, matryoshka_dims=matryoshka_dims)
# Evaluators for validation
evaluators = []
for dim in matryoshka_dims:
evaluators.append(
TripletEvaluator(
anchors=eval_subset["anchor"],
positives=eval_subset["positive"],
negatives=eval_subset["negative"],
name=f"dev-{dim}",
truncate_dim=dim,
)
)
dev_evaluator = SequentialEvaluator(evaluators, main_score_function=lambda scores: scores[0])
dev_evaluator(model)
args = SentenceTransformerTrainingArguments(
# --- Output & Identity ---
output_dir=output_dir, # Where the model and checkpoints are saved
seed=42, # Ensures results are reproducible (shuffling/init)
# --- Epochs & Batching (L40S Optimized) ---
num_train_epochs=2, # Increased to 2 for better convergence
per_device_train_batch_size=train_batch_size, # High batch size to saturate the L40S 48GB VRAM
per_device_eval_batch_size=train_batch_size, # Matching eval batch size for speed
gradient_accumulation_steps=2, # Effective batch size = 128 (64 * 2)
# --- Optimization & Precision ---
bf16=True, # Set to True for L40S; faster and more stable than FP16
fp16=False, # Disabled in favor of BF16
learning_rate=2e-5, # Standard "safe" learning rate for Transformers
lr_scheduler_type="linear", # Gently reduces learning rate to 0 over training
warmup_ratio=0.1, # Ramps up LR for the first 10% of steps to prevent spikes
weight_decay=0.01, # Regularization to prevent overfitting
# --- Data Handling ---
batch_sampler=BatchSamplers.NO_DUPLICATES, # Essential for MNR/Matryoshka loss to avoid bad negatives
dataloader_num_workers=8, # Use 8 CPU cores to keep the GPU fed with data
# --- Evaluation & Saving (Safety) ---
eval_strategy="steps", # Evaluate every X steps
eval_steps=6000, # Increased from 10; L40S processes data very fast
save_strategy="steps", # Save checkpoints every X steps
save_steps=6000, # Usually matches eval_steps
save_total_limit=2, # Keep only the top 3 checkpoints to save disk space
# --- Tracking & Debugging ---
report_to="tensorboard", # Send live metrics to TensorBoard
logging_steps=200, # Print/Log stats every 200 steps (prevents messy logs)
logging_dir="arabvert02-matryoshka/runs", # Specific folder for TensorBoard event files
)
trainer = SentenceTransformerTrainer(
model=model,
args=args,
train_dataset=train_dataset,
eval_dataset=eval_dataset,
loss=train_loss,
evaluator=dev_evaluator,
)
trainer.train()
# Save final model
final_output_dir = f"{output_dir}/final"
model.save(final_output_dir)
print("model saved successfully")
# Test evaluation
evaluators = []
for dim in matryoshka_dims:
evaluators.append(
TripletEvaluator(
anchors=test_dataset["train"]["anchor"],
positives=test_dataset["train"]["positive"],
negatives=test_dataset["train"]["negative"],
name=f"test-{dim}",
truncate_dim=dim,
)
)
test_evaluator = SequentialEvaluator(evaluators)
test_evaluator(model)
|