Upload train.py with huggingface_hub

train.py

@@ -0,0 +1,54 @@
+"""
+Ult1-Coding Fine-Tuning Script
+Fine-tune on any code dataset.
+Requires GPU with ~10 GB VRAM.
+
+Usage:
+    python train.py --dataset bigcode/the-stack-dedup --subset data/python
+"""
+import torch, argparse, os
+from transformers import AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer, DataCollatorForSeq2Seq
+from peft import LoraConfig, get_peft_model, TaskType
+from datasets import load_dataset
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model", default="teolm30/Ult1-coding")
+parser.add_argument("--dataset", default="code_search_net")
+parser.add_argument("--subset", default="python")
+parser.add_argument("--lr", type=float, default=2e-4)
+parser.add_argument("--epochs", type=int, default=1)
+parser.add_argument("--max_length", type=int, default=1024)
+parser.add_argument("--output", default="./ult1_coding_finetuned")
+args = parser.parse_args()
+
+os.makedirs(args.output, exist_ok=True)
+
+model = AutoModelForCausalLM.from_pretrained(args.model, torch_dtype=torch.bfloat16, device_map="auto")
+tokenizer = AutoTokenizer.from_pretrained(args.model)
+if tokenizer.pad_token is None:
+    tokenizer.pad_token = tokenizer.eos_token
+
+lora_config = LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj"], task_type=TaskType.CAUSAL_LM)
+model = get_peft_model(model, lora_config)
+model.print_trainable_parameters()
+
+dataset = load_dataset(args.dataset, split="train", streaming=True).take(5000)
+def fmt(ex):
+    code = ex.get("code") or ex.get("content") or ex.get("func_code_string") or str(ex)
+    return {"text": f"<|im_start|>user\nWrite code:\n<|im_end|>\n<|im_start|>assistant\n{code}<|im_end|>"}
+dataset = dataset.map(fmt)
+def tok(exs):
+    return tokenizer(exs["text"], truncation=True, max_length=args.max_length, padding="max_length")
+dataset = dataset.map(tok, remove_columns=[c for c in dataset.column_names if c != "text"], batched=True)
+
+args_t = TrainingArguments(
+    output_dir=args.output, per_device_train_batch_size=2, gradient_accumulation_steps=8,
+    num_train_epochs=args.epochs, learning_rate=args.lr, logging_steps=10,
+    save_strategy="epoch", bf16=True, report_to="none",
+)
+trainer = Trainer(model=model, args=args_t, train_dataset=dataset,
+                  data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8))
+trainer.train()
+model.save_pretrained(args.output)
+tokenizer.save_pretrained(args.output)
+print(f"Saved to {args.output}")