Fox-AI-by-teolm30/Ult1-coding

1"""Ult1-Coding GPU fine-tuning script.2Usage: python train_coding.py (requires 8+ GB VRAM GPU)3"""4import json, torch, warnings5from datasets import Dataset6from transformers import (7    AutoModelForCausalLM, AutoTokenizer, TrainingArguments, Trainer8)9from peft import LoraConfig, get_peft_model, TaskType10from huggingface_hub import HfApi11warnings.filterwarnings("ignore")12 13import os14HF_TOKEN = os.getenv("HF_TOKEN", "")15MODEL_ID = "Qwen/Qwen2.5-3B-Instruct"16REPO_ID = "teolm30/Ult1-coding"17 18def format_example(example):19    return {20        "text": f"<|im_start|>user\n{example['instruction']}<|im_end|>\n<|im_start|>assistant\n{example['response']}<|im_end|>\n"21    }22 23print("Loading tokenizer...")24tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)25tokenizer.pad_token = tokenizer.eos_token26 27print("Loading training data from HF...")28api = HfApi()29data_path = api.hf_hub_download(repo_id=REPO_ID, filename="training_data.json", token=HF_TOKEN)30with open(data_path, "r") as f:31    raw_data = json.load(f)32 33formatted = [format_example(ex) for ex in raw_data]34dataset = Dataset.from_list(formatted)35 36def tokenize_fn(examples):37    result = tokenizer(38        examples["text"], truncation=True, max_length=512,39        padding="max_length", return_tensors=None40    )41    result["labels"] = result["input_ids"].copy()42    return result43 44dataset = dataset.map(tokenize_fn, remove_columns=["text"], batched=True)45dataset = dataset.train_test_split(test_size=0.1)46 47print(f"Training samples: {len(dataset['train'])}, Validation: {len(dataset['test'])}")48 49print("Loading model...")50model = AutoModelForCausalLM.from_pretrained(51    MODEL_ID, torch_dtype=torch.bfloat16,52    device_map="auto", attn_implementation="flash_attention_2",53)54 55model.train()56lora_config = LoraConfig(57    r=16, lora_alpha=32, lora_dropout=0.05,58    target_modules=["q_proj", "k_proj", "v_proj", "o_proj",59                    "gate_proj", "up_proj", "down_proj"],60    bias="none", task_type=TaskType.CAUSAL_LM,61)62model = get_peft_model(model, lora_config)63model.print_trainable_parameters()64 65training_args = TrainingArguments(66    output_dir="./ult1_coding_trained",67    per_device_train_batch_size=1,68    gradient_accumulation_steps=4,69    learning_rate=3e-4,70    num_train_epochs=3,71    logging_steps=1,72    save_strategy="epoch",73    evaluation_strategy="epoch",74    bf16=True,75    gradient_checkpointing=True,76    logging_dir="./logs",77)78 79trainer = Trainer(80    model=model, args=training_args,81    train_dataset=dataset["train"],82    eval_dataset=dataset["test"],83)84 85trainer.train()86model.save_pretrained("./ult1_coding_trained/final", safe_serialization=True)87tokenizer.save_pretrained("./ult1_coding_trained/final")88 89api.upload_folder(90    folder_path="./ult1_coding_trained/final",91    repo_id=REPO_ID, token=HF_TOKEN,92    commit_message="Ult1-Coding trained adapter - GPU fine-tuned on coding data",93    ignore_patterns=["*.pt", "checkpoint-*"],94)95print(f"Trained model uploaded to https://huggingface.co/{REPO_ID}")96