HeeBive commited on
Commit
8b8bde9
·
verified ·
1 Parent(s): 8093f95

Upload train.sh with huggingface_hub

Browse files
Files changed (1) hide show
  1. train.sh +44 -0
train.sh ADDED
@@ -0,0 +1,44 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ #!/bin/bash
2
+ set -e
3
+ pip install torch transformers peft trl datasets bitsandbytes accelerate sentencepiece protobuf huggingface_hub -q 2>&1 | tail -1
4
+
5
+ python3 << 'PYEOF' > training.log 2>&1
6
+ import os, torch, shutil, json
7
+ os.environ["TORCH_COMPILE_DISABLE"] = "1"
8
+ from huggingface_hub import hf_hub_download
9
+ from datasets import Dataset
10
+ from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig, TrainingArguments
11
+ from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
12
+ from trl import SFTTrainer
13
+
14
+ MODEL = "mistralai/Mixtral-8x7B-Instruct-v0.1"
15
+ OUT = "./adapter"
16
+ TOKEN = os.environ.get("HF_TOKEN", "")
17
+
18
+ print("DL data...")
19
+ p = hf_hub_download("HeeBive/ZeroSec-7B", "training_data.jsonl", repo_type="model", token=TOKEN)
20
+ raw = [json.loads(l) for l in open(p) if l.strip()]
21
+ print(f"{len(raw)} samples")
22
+ def f(e): return {"text": f"<|im_start|>user\n{e['instruction']}<|im_end|>\n<|im_start|>assistant\n{e['response']}<|im_end|>"}
23
+ ds = Dataset.from_list(raw).map(f)
24
+ tok = AutoTokenizer.from_pretrained(MODEL)
25
+ tok.pad_token = tok.eos_token
26
+ ds = ds.map(lambda x: tok(x["text"], truncation=True, max_length=512), batched=True, remove_columns=ds.column_names)
27
+
28
+ print("Load model...")
29
+ bnb = BitsAndBytesConfig(load_in_4bit=True, bnb_4bit_use_double_quant=True, bnb_4bit_quant_type="nf4", bnb_4bit_compute_dtype=torch.bfloat16)
30
+ m = AutoModelForCausalLM.from_pretrained(MODEL, quantization_config=bnb, device_map="auto", torch_dtype=torch.bfloat16)
31
+ m = prepare_model_for_kbit_training(m)
32
+ m.gradient_checkpointing_enable()
33
+ 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"))
34
+
35
+ SFTTrainer(model=m, processing_class=tok, train_dataset=ds,
36
+ args=TrainingArguments(output_dir=OUT, per_device_train_batch_size=1, gradient_accumulation_steps=8,
37
+ num_train_epochs=1, learning_rate=1e-4, bf16=True, logging_steps=5, optim="adamw_8bit", report_to="none")
38
+ ).train()
39
+ m.save_pretrained(OUT); tok.save_pretrained(OUT)
40
+ if TOKEN:
41
+ from huggingface_hub import HfApi
42
+ HfApi().upload_folder(folder_path=OUT, repo_id=os.environ.get("HF_REPO_ID","HeeBive/ZeroSec"), repo_type="model", token=TOKEN)
43
+ print(f"✅ Uploaded")
44
+ PYEOF