Spaces:
Sleeping
Sleeping
| --- | |
| # ✅ train.py — Trains StarCoder 7B on your dataset | |
| from datasets import load_dataset | |
| from transformers import AutoTokenizer, AutoModelForCausalLM, TrainingArguments, Trainer, DataCollatorForLanguageModeling | |
| import os | |
| import sys | |
| print("🔥 Training script started...", file=sys.stderr) | |
| # === CONFIG === | |
| DATASET_PATH = "python_ai_dataset.jsonl" # Must exist in Space root | |
| MODEL_ID = "bigcode/starcoderbase-7b" | |
| OUTPUT_DIR = "train_output" | |
| # === Load Dataset === | |
| dataset = load_dataset("json", data_files=DATASET_PATH, split="train") | |
| # === Load Tokenizer and Model === | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained(MODEL_ID, trust_remote_code=True) | |
| # === Preprocessing === | |
| def tokenize(example): | |
| return tokenizer(example["prompt"] + "\n" + example["completion"], truncation=True, max_length=512) | |
| tokenized_dataset = dataset.map(tokenize, remove_columns=["prompt", "completion"]) | |
| data_collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False) | |
| # === Training Args === | |
| training_args = TrainingArguments( | |
| output_dir=OUTPUT_DIR, | |
| overwrite_output_dir=True, | |
| per_device_train_batch_size=1, | |
| gradient_accumulation_steps=4, | |
| num_train_epochs=2, | |
| logging_dir="./logs", | |
| logging_steps=10, | |
| save_strategy="epoch", | |
| save_total_limit=2, | |
| fp16=True, | |
| bf16=False, | |
| report_to="none", | |
| ) | |
| # === Train === | |
| trainer = Trainer( | |
| model=model, | |
| args=training_args, | |
| train_dataset=tokenized_dataset, | |
| tokenizer=tokenizer, | |
| data_collator=data_collator | |
| ) | |
| trainer.train() | |
| # === Save === | |
| trainer.save_model(OUTPUT_DIR) | |
| tokenizer.save_pretrained(OUTPUT_DIR) |