File size: 2,855 Bytes
58d9159 ecd6fbd 58d9159 | 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 | import math
import torch
from tokenizers import Tokenizer
from transformers import PreTrainedTokenizerFast, get_cosine_schedule_with_warmup
from training import PreTrainer
from tynerox.modeling import TyneRoxModel, TyneRoxConfig
from dataset.pre_train import create_train_dataloader
if __name__ == "__main__":
# 1 - Carrega o tokenizador
tokenizer = Tokenizer.from_file("tokenizer/tokens-bpe-36k.json")
tokenizer = PreTrainedTokenizerFast(
tokenizer_object=tokenizer,
unk_token="[UNK]",
pad_token="<|endoftext|>",
eos_token="<|endoftext|>",
)
tokenizer.save_pretrained(f"../")
# 2 Inicia a configuração e o modelo
config = TyneRoxConfig(
vocab_size=tokenizer.vocab_size,
pad_token_id=tokenizer.pad_token_id,
)
model = TyneRoxModel(config)
model.to("cuda")
# 3 - Carrega o dataset de treinamento
folder_path = "bobboyms/subset-Itau-Unibanco-aroeira-1B-tokens"
dataloader = create_train_dataloader(
folder_path,
tokenizer,
batch_size=5,
max_length=1024,
drop_last=True,
num_workers=10
)
# 4 - Criando o optmizer
model = torch.compile(model)
optimizer = torch.optim.AdamW(
model.parameters(),
lr=0.000461, # Mantenha a LR inicial ou ajuste ligeiramente (ex: 3e-4)
weight_decay=0.1
)
# 5 - Configura o warmup
epochs = 1
batch_size = 40
size_dataset = 2_883_231
warmup_ratio = 0.05
num_training_steps = len(dataloader) * epochs
num_warmup_steps = math.floor(num_training_steps * warmup_ratio)
# 6. Scheduler
scheduler = get_cosine_schedule_with_warmup(
optimizer,
num_warmup_steps=num_warmup_steps,
num_training_steps=num_training_steps,
)
sample_prompts = [
"Olá, como vai você? ",
"Quando a manhã chegou, Iracema ainda estava ali, debruçada, como uma borboleta que ",
"Não, respondeu; na verdade, estou com medo ",
"O resultado representa uma desaceleração ",
"No vídeo, é possível ver ",
"Essa receita de torta de frango ",
"Durante o primeiro mandato ",
"Os donos de cães "
]
logger_config = {
"tracking_uri": "http://127.0.0.1:5000",
"experiment": "Pre training LLM",
"model_name": "Pre training LLM (Long Context)"
}
trainer = PreTrainer(
model=model,
optimizer=optimizer,
scheduler=scheduler,
tokenizer=tokenizer,
train_loader=dataloader,
test_loader=None,
logger_config=logger_config,
use_amp=True
)
trainer.train(num_epochs=epochs,sample_prompts=sample_prompts)
# 7 - Salva as configurações do modelo para enviar para o hugginfaces
model.save_pretrained(f"../")
|