Upload train.py with huggingface_hub

train.py

@@ -1,54 +1,95 @@
+"""Ult1-Coding GPU fine-tuning script.
+Usage: python train_coding.py (requires 8+ GB VRAM GPU)
 """
-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
+import json, torch, warnings
+from datasets import Dataset
+from transformers import (
+    AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer
+)
 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)
+from huggingface_hub import HfApi
+warnings.filterwarnings("ignore")
+
+import os
+HF_TOKEN = os.getenv("HF_TOKEN", "")
+MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"
+REPO_ID = "teolm30/Ult1-coding"
+
+def format_example(example):
+    return {
+        "text": f"<|im_start|>user\n{example['instruction']}<|im_end|>\n<|im_start|>assistant\n{example['response']}<|im_end|>\n"
+    }
+
+print("Loading tokenizer...")
+tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+tokenizer.pad_token = tokenizer.eos_token
+
+print("Loading training data from HF...")
+api = HfApi()
+data_path = api.hf_hub_download(repo_id=REPO_ID, filename="training_data.json", token=HF_TOKEN)
+with open(data_path, "r") as f:
+    raw_data = json.load(f)
+
+formatted = [format_example(ex) for ex in raw_data]
+dataset = Dataset.from_list(formatted)
+
+def tokenize_fn(examples):
+    result = tokenizer(
+        examples["text"], truncation=True, max_length=512,
+        padding="max_length", return_tensors=None
+    )
+    result["labels"] = result["input_ids"].copy()
+    return result
+
+dataset = dataset.map(tokenize_fn, remove_columns=["text"], batched=True)
+dataset = dataset.train_test_split(test_size=0.1)
+
+print(f"Training samples: {len(dataset['train'])}, Validation: {len(dataset['test'])}")
+
+print("Loading model...")
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID, torch_dtype=torch.bfloat16,
+    device_map="auto", attn_implementation="flash_attention_2",
+)
+
+model.train()
+lora_config = LoraConfig(
+    r=16, lora_alpha=32, lora_dropout=0.05,
+    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
+                    "gate_proj", "up_proj", "down_proj"],
+    bias="none", 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",
+training_args = TrainingArguments(
+    output_dir="./ult1_coding_trained",
+    per_device_train_batch_size=1,
+    gradient_accumulation_steps=4,
+    learning_rate=3e-4,
+    num_train_epochs=3,
+    logging_steps=1,
+    save_strategy="epoch",
+    evaluation_strategy="epoch",
+    bf16=True,
+    gradient_checkpointing=True,
+    logging_dir="./logs",
 )
-trainer = Trainer(model=model, args=args_t, train_dataset=dataset,
-                  data_collator=DataCollatorForSeq2Seq(tokenizer, pad_to_multiple_of=8))
+
+trainer = Trainer(
+    model=model, args=training_args,
+    train_dataset=dataset["train"],
+    eval_dataset=dataset["test"],
+)
+
 trainer.train()
-model.save_pretrained(args.output)
-tokenizer.save_pretrained(args.output)
-print(f"Saved to {args.output}")
+model.save_pretrained("./ult1_coding_trained/final", safe_serialization=True)
+tokenizer.save_pretrained("./ult1_coding_trained/final")
+
+api.upload_folder(
+    folder_path="./ult1_coding_trained/final",
+    repo_id=REPO_ID, token=HF_TOKEN,
+    commit_message="Ult1-Coding trained adapter - GPU fine-tuned on coding data",
+    ignore_patterns=["*.pt", "checkpoint-*"],
+)
+print(f"Trained model uploaded to https://huggingface.co/{REPO_ID}")