ByteBandits/train.py at main · RSAgr/ByteBandits · GitHub

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
import os
import torch
from transformers import (
    AutoModelForCausalLM,
    AutoTokenizer,
    Trainer,
    TrainingArguments,
    DataCollatorForLanguageModeling,
)
from datasets import load_dataset, Dataset
from peft import get_peft_model, LoraConfig, TaskType
import json

# Load model and tokenizer
# model_name = "deepseek-ai/deepseek-coder-1.3b-base"
model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
tokenizer = AutoTokenizer.from_pretrained(model_name)
tokenizer.pad_token = tokenizer.eos_token
model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float32)

# Enable LoRA
# peft_config = LoraConfig(
#     task_type=TaskType.CAUSAL_LM,
#     inference_mode=False,
#     r=8,
#     lora_alpha=16,
#     lora_dropout=0.05,
# )
# model = get_peft_model(model, peft_config)

from peft import PeftModel, PeftConfig, get_peft_model, LoraConfig, TaskType

peft_path = "./trainedModel"

# if os.path.exists(peft_path):
#     # ✅ Load existing LoRA adapter (continue training)
#     print("🔁 Loading existing LoRA weights...")
#     model = PeftModel.from_pretrained(model, peft_path)
# else:
#     # 🆕 No existing LoRA adapter found — create a new one
#     print("🆕 No previous LoRA weights found. Initializing new adapter...")
#     peft_config = LoraConfig(
#         task_type=TaskType.CAUSAL_LM,
#         inference_mode=False,
#         r=8,
#         lora_alpha=16,
#         lora_dropout=0.05,
#     )
#     model = get_peft_model(model, peft_config)

if os.path.exists(os.path.join(peft_path, "adapter_config.json")):
    print("🔁 Loading existing LoRA weights...")
    model = PeftModel.from_pretrained(model, peft_path)
else:
    print("🆕 No previous LoRA weights found. Initializing new adapter...")
    peft_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        inference_mode=False,
        r=8,
        lora_alpha=16,
        lora_dropout=0.05,
    )
    model = get_peft_model(model, peft_config)


model.train()
model.requires_grad_(True)

# Load dataset
def load_jsonl_dataset(path):
    with open(path, "r", encoding="utf-8") as f:
        lines = [json.loads(line) for line in f]
    return Dataset.from_list([
        {"text": f"Input:\n{ex['instruction']}\nOutput:\n{ex['output']}"}
        for ex in lines
    ])

dataset = load_jsonl_dataset(r"data\samples.jsonl")
dataset = dataset.select(range(0,10))

# Tokenize and add labels
def tokenize(example):
    tokenized = tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)
    #tokenized = tokenizer(example["text"], truncation=True, padding="max_length", max_length=512)
    tokenized["labels"] = tokenized["input_ids"].copy()  # Required for loss computation
    return tokenized

tokenized_dataset = dataset.map(tokenize, batched=True)

print(tokenized_dataset[0])


# Training config
training_args = TrainingArguments(
    output_dir="./trainedModel",
    per_device_train_batch_size=1,
    num_train_epochs=10,
    save_steps=100,
    save_total_limit=2,
    learning_rate=1e-4,
    logging_dir="./logs",
    logging_steps=10,
    gradient_checkpointing=False,  # Disable if no GPU
    fp16=False,
    push_to_hub=False,
)

# Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=DataCollatorForLanguageModeling(tokenizer, mlm=False),
)
logging_steps=1,
logging_first_step=True,

trainer.train()