| #!/bin/bash |
| set -e |
| pip install torch transformers peft trl datasets bitsandbytes accelerate sentencepiece protobuf huggingface_hub -q 2>&1 | tail -1 |
|
|
| TOKEN="${HF_TOKEN:?Set HF_TOKEN first}" |
| REPO="${HF_REPO_ID:-HeeBive/ZeroSec}" |
|
|
| python3 << 'PYEOF' > training.log 2>&1 |
| import os, torch, shutil, json |
| os.environ["TORCH_COMPILE_DISABLE"] = "1" |
| TOKEN = os.environ.get("HF_TOKEN") |
| REPO = os.environ.get("HF_REPO_ID", "HeeBive/ZeroSec") |
| from huggingface_hub import hf_hub_download, HfApi |
| from datasets import Dataset |
| from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments |
| from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training |
| from trl import SFTTrainer |
|
|
| MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1" |
| OUT = "./adapter" |
|
|
| print("DL data...") |
| p = hf_hub_download("HeeBive/ZeroSec-7B", "training_data.jsonl", repo_type="model", token=TOKEN) |
| raw = [json.loads(l) for l in open(p) if l.strip()] |
| print(f"{len(raw)} samples") |
| def f(e): return {"text": f"<|im_start|>user\n{e['instruction']}<|im_end|>\n<|im_start|>assistant\n{e['response']}<|im_end|>"} |
| ds = Dataset.from_list(raw).map(f) |
| tok = AutoTokenizer.from_pretrained(MODEL) |
| tok.pad_token = tok.eos_token |
| ds = ds.map(lambda x: tok(x["text"], truncation=True, max_length=512), batched=True, remove_columns=ds.column_names) |
|
|
| print("Load model...") |
| bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16) |
| m = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb, device_map="auto", torch_dtype=torch.bfloat16) |
| m = prepare_model_for_kbit_training(m) |
| m.gradient_checkpointing_enable() |
| m = get_peft_model(m, LoraConfig(r=16, lora_alpha=32, target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"], task_type="CAUSAL_LM", bias="none")) |
| SFTTrainer(model=m, processing_class=tok, train_dataset=ds, |
| args=TrainingArguments(output_dir=OUT, per_device_train_batch_size=1, gradient_accumulation_steps=8, num_train_epochs=1, learning_rate=1e-4, bf16=True, logging_steps=5, optim="adamw_8bit", report_to="none")).train() |
| m.save_pretrained(OUT); tok.save_pretrained(OUT) |
| print("Uploading...") |
| HfApi(token=TOKEN).upload_folder(folder_path=OUT, repo_id=REPO, repo_type="model", token=TOKEN) |
| print(f"✅ {REPO}") |
| PYEOF |
|
|
