File size: 4,633 Bytes
6f52d03 | 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 | #!/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.")
|