144 lines
4.5 KiB
Python
144 lines
4.5 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
ChatPBC V3.3 Training Script
|
|
Train this model on a GPU-enabled environment.
|
|
|
|
Usage:
|
|
python train_v4.py
|
|
|
|
Requirements:
|
|
- GPU with 16GB+ VRAM (A10G, T4, L4, A100, etc.)
|
|
- transformers, peft, bitsandbytes, accelerate, datasets
|
|
"""
|
|
|
|
import torch
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments, Trainer
|
|
from peft import LoraConfig, get_peft_model, TaskType
|
|
from datasets import load_dataset
|
|
|
|
HF_TOKEN = "YOUR_HF_TOKEN_HERE"
|
|
BASE_MODEL = "mistralai/Mistral-7B-v0.1"
|
|
DATASET_PATH = "chatpbc1/chatpbc-business-consulting-dataset"
|
|
OUTPUT_REPO = "chatpbc1/chatpbc-v33"
|
|
|
|
def reformat_to_mistral(example, max_target_length=400):
|
|
text = example["text"]
|
|
instruction = ""
|
|
input_text = ""
|
|
response = ""
|
|
|
|
if "### Instruction:" in text and "### Response:" in text:
|
|
parts = text.split("### Response:")
|
|
pre_response = parts[0]
|
|
response = parts[1].strip() if len(parts) > 1 else ""
|
|
|
|
if "### Input:" in pre_response:
|
|
inst_parts = pre_response.split("### Input:")
|
|
instruction = inst_parts[0].replace("### Instruction:", "").strip()
|
|
input_text = inst_parts[1].strip()
|
|
else:
|
|
instruction = pre_response.replace("### Instruction:", "").strip()
|
|
else:
|
|
instruction = text
|
|
|
|
words = response.split()
|
|
if len(words) > max_target_length:
|
|
response = " ".join(words[:max_target_length])
|
|
|
|
if input_text:
|
|
prompt = f"<s>[INST] {instruction}\\n\\nContext: {input_text} [/INST]"
|
|
else:
|
|
prompt = f"<s>[INST] {instruction} [/INST]"
|
|
|
|
return {"text": prompt + response + "</s>"}
|
|
|
|
def tokenize_fn(example):
|
|
return tokenizer(example["text"], truncation=True, max_length=2048, padding=False)
|
|
|
|
# Setup
|
|
from huggingface_hub import login
|
|
login(token=HF_TOKEN)
|
|
|
|
tokenizer = AutoTokenizer.from_pretrained(BASE_MODEL, token=HF_TOKEN, use_fast=True)
|
|
tokenizer.pad_token = tokenizer.eos_token
|
|
|
|
bnb_config = BitsAndBytesConfig(
|
|
load_in_4bit=True,
|
|
bnb_4bit_quant_type="nf4",
|
|
bnb_4bit_compute_dtype=torch.float16,
|
|
bnb_4bit_use_double_quant=True,
|
|
)
|
|
|
|
model = AutoModelForCausalLM.from_pretrained(
|
|
BASE_MODEL, quantization_config=bnb_config, device_map="auto",
|
|
torch_dtype=torch.float16, token=HF_TOKEN,
|
|
)
|
|
|
|
from peft import prepare_model_for_kbit_training
|
|
model = prepare_model_for_kbit_training(model)
|
|
model.config.use_cache = False
|
|
|
|
lora_config = LoraConfig(
|
|
r=16, lora_alpha=32,
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
lora_dropout=0.05, bias="none", task_type=TaskType.CAUSAL_LM,
|
|
)
|
|
|
|
model = get_peft_model(model, lora_config)
|
|
model.print_trainable_parameters()
|
|
|
|
# Data
|
|
dataset = load_dataset(DATASET_PATH, split="train")
|
|
sample_size = min(len(dataset), 10000)
|
|
sampled = dataset.shuffle(seed=42).select(range(sample_size))
|
|
formatted = sampled.map(lambda x: reformat_to_mistral(x, 1000), remove_columns=sampled.column_names)
|
|
tokenized = formatted.map(tokenize_fn, remove_columns=formatted.column_names, num_proc=2)
|
|
|
|
# Train
|
|
class SimpleCollator:
|
|
def __call__(self, features):
|
|
input_ids = [f["input_ids"] for f in features]
|
|
max_len = max(len(x) for x in input_ids)
|
|
padded = [x + [tokenizer.pad_token_id] * (max_len - len(x)) for x in input_ids]
|
|
attention_mask = [[1] * len(x) + [0] * (max_len - len(x)) for x in input_ids]
|
|
labels = padded.copy()
|
|
return {
|
|
"input_ids": torch.tensor(padded, dtype=torch.long),
|
|
"attention_mask": torch.tensor(attention_mask, dtype=torch.long),
|
|
"labels": torch.tensor(labels, dtype=torch.long),
|
|
}
|
|
|
|
training_args = TrainingArguments(
|
|
output_dir="./chatpbc-v4-output",
|
|
num_train_epochs=3,
|
|
max_steps=1000,
|
|
per_device_train_batch_size=2,
|
|
gradient_accumulation_steps=8,
|
|
learning_rate=2e-4,
|
|
weight_decay=0.01,
|
|
warmup_ratio=0.03,
|
|
lr_scheduler_type="cosine",
|
|
fp16=True,
|
|
gradient_checkpointing=True,
|
|
logging_steps=50,
|
|
save_steps=500,
|
|
save_strategy="steps",
|
|
evaluation_strategy="no",
|
|
push_to_hub=True,
|
|
hub_model_id=OUTPUT_REPO,
|
|
hub_token=HF_TOKEN,
|
|
report_to="none",
|
|
optim="paged_adamw_8bit",
|
|
max_grad_norm=0.3,
|
|
)
|
|
|
|
trainer = Trainer(
|
|
model=model, args=training_args, train_dataset=tokenized,
|
|
data_collator=SimpleCollator(), tokenizer=tokenizer,
|
|
)
|
|
|
|
trainer.train()
|
|
model.push_to_hub(OUTPUT_REPO, token=HF_TOKEN)
|
|
tokenizer.push_to_hub(OUTPUT_REPO, token=HF_TOKEN)
|
|
print("Training complete! Model uploaded to HF Hub.")
|