File size: 3,448 Bytes
ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 f5d9b2e ca82b75 7860651 ca82b75 |
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 |
import argparse, os, sys
from typing import List
from datasets import load_dataset
from transformers import (
AutoTokenizer, AutoModelForCausalLM,
DataCollatorForLanguageModeling, TrainingArguments, Trainer
)
def parse_args():
p = argparse.ArgumentParser()
p.add_argument("--dataset", required=True, help="JSON/JSONL (.jsonl or .jsonl.gz)")
p.add_argument("--output", default="trained_model")
p.add_argument("--model_name", default="distilgpt2")
p.add_argument("--epochs", type=float, default=0.5)
p.add_argument("--batch_size", type=int, default=2)
p.add_argument("--block_size", type=int, default=256)
p.add_argument("--learning_rate", type=float, default=5e-5)
# quick mode:
p.add_argument("--quick", type=int, default=0) # 1 => tiny model + fast
p.add_argument("--max_steps", type=int, default=0) # >0 overrides epochs
p.add_argument("--subset", type=int, default=0) # use first N rows
return p.parse_args()
def main():
a = parse_args()
if a.quick:
a.model_name = "sshleifer/tiny-gpt2" # ultra-tiny, very fast
if a.max_steps <= 0: a.max_steps = 8
if a.subset <= 0: a.subset = 32
a.epochs = 1.0
print(f"📥 Loading dataset: {a.dataset}", flush=True)
ds = load_dataset("json", data_files=a.dataset, split="train")
cols = ds.column_names
print("🧾 Columns:", cols, flush=True)
if a.subset and a.subset > 0:
ds = ds.select(range(min(a.subset, len(ds))))
print(f"✂ Using subset: {len(ds)} rows", flush=True)
tok = AutoTokenizer.from_pretrained(a.model_name)
if tok.pad_token is None:
tok.pad_token = tok.eos_token
model = AutoModelForCausalLM.from_pretrained(a.model_name)
def build_texts(batch) -> List[str]:
if "text" in batch:
return [str(t) for t in batch["text"]]
if "prompt" in batch and "completion" in batch:
return [f"{str(p).rstrip()}\n{str(c)}" for p, c in zip(batch["prompt"], batch["completion"])]
raise ValueError("Dataset must contain 'text' OR both 'prompt' and 'completion'.")
def tokenize(batch):
texts = build_texts(batch)
return tok(texts, padding="max_length", truncation=True, max_length=a.block_size)
print("🔁 Tokenizing…", flush=True)
tokds = ds.map(tokenize, batched=True, remove_columns=cols)
collator = DataCollatorForLanguageModeling(tokenizer=tok, mlm=False)
print("⚙ Trainer…", flush=True)
args = TrainingArguments(
output_dir=a.output,
overwrite_output_dir=True,
per_device_train_batch_size=a.batch_size,
num_train_epochs=a.epochs if a.max_steps == 0 else 1,
learning_rate=a.learning_rate,
logging_steps=1,
save_steps=50,
save_total_limit=1,
report_to=[],
fp16=False,
max_steps=a.max_steps if a.max_steps > 0 else -1,
)
trainer = Trainer(model=model, args=args, train_dataset=tokds, tokenizer=tok, data_collator=collator)
print("🚀 Training…", flush=True)
trainer.train()
print(f"💾 Saving to {a.output}", flush=True)
os.makedirs(a.output, exist_ok=True)
trainer.save_model(a.output)
tok.save_pretrained(a.output)
print("✅ Done.", flush=True)
if __name__ == "__main__":
try:
main()
except Exception as e:
print(f"❌ Training failed: {e}", flush=True)
sys.exit(1) |