File size: 2,344 Bytes
8b8bde9 c1404ed 8b8bde9 c1404ed 8b8bde9 c1404ed 8b8bde9 c1404ed 8b8bde9 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 | #!/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
|