Spaces:
Sleeping
Sleeping
| import argparse | |
| from datasets import load_dataset | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling | |
| parser = argparse.ArgumentParser() | |
| parser.add_argument("--dataset", required=True) | |
| parser.add_argument("--output", default="trained_model") | |
| args = parser.parse_args() | |
| print("📊 Loading dataset...") | |
| dataset = load_dataset("json", data_files=args.dataset, split="train") | |
| print("🧠 Loading model and tokenizer...") | |
| tokenizer = AutoTokenizer.from_pretrained("distilgpt2") | |
| tokenizer.pad_token = tokenizer.eos_token | |
| model = AutoModelForCausalLM.from_pretrained("distilgpt2") | |
| # ✅ Clean, batch-safe tokenize | |
| def tokenize(batch): | |
| full_texts = [str(p) + str(c) for p, c in zip(batch["prompt"], batch["completion"])] | |
| return tokenizer(full_texts, padding="max_length", truncation=True, max_length=256) | |
| print("🔁 Tokenizing...") | |
| tokenized = dataset.map(tokenize, batched=True) | |
| print("📦 Setting up trainer...") | |
| data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) | |
| training_args = TrainingArguments( | |
| output_dir=args.output, | |
| per_device_train_batch_size=2, | |
| num_train_epochs=1, | |
| logging_steps=1, | |
| save_steps=5, | |
| save_total_limit=1, | |
| report_to=[] | |
| ) | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=tokenized, | |
| tokenizer=tokenizer, | |
| data_collator=data_collator, | |
| ) | |
| print("🚀 Starting training...") | |
| trainer.train() | |
| trainer.save_model(args.output) | |
| print("✅ Done.") |