Spaces:
Running
Running
File size: 4,229 Bytes
aeae5cb | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """
Unsloth local training script for InfraMind (Aegis-Swarm).
Requirements: pip install "unsloth[colab-new] @ git+https://github.com/unslothai/unsloth.git" trl peft
"""
import os
import torch
from datasets import load_dataset
from trl import SFTTrainer
from transformers import TrainingArguments
# Try importing Unsloth. If it fails, fallback to standard transformers or mock for demo
try:
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
HAS_UNSLOTH = True
except ImportError:
HAS_UNSLOTH = False
print("Warning: Unsloth not found. Install it for 2x faster training.")
def format_chat_template(example):
# This aligns the JSONL 'conversations' with the ShareGPT schema
return {"messages": example["conversations"]}
def main():
if not os.path.exists("data/inframind_sft.jsonl"):
print("Error: data/inframind_sft.jsonl not found. Run scripts/generate_dataset.py first.")
return
print("Loading InfraMind SFT Dataset...")
dataset = load_dataset("json", data_files={"train": "data/inframind_sft.jsonl"}, split="train")
# In a real environment, we would map the conversations to a standard format
# For now, we assume standard ShareGPT format from the generator
max_seq_length = 4096
model_name = "unsloth/llama-3-8b-Instruct-bnb-4bit"
if HAS_UNSLOTH:
print(f"Initializing Unsloth Model: {model_name}")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=model_name,
max_seq_length=max_seq_length,
dtype=None,
load_in_4bit=True,
)
model = FastLanguageModel.get_peft_model(
model,
r=16,
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
lora_alpha=16,
lora_dropout=0,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=42,
use_rslora=False,
loftq_config=None,
)
tokenizer = get_chat_template(
tokenizer,
chat_template="llama-3",
mapping={"role": "role", "content": "content", "user": "user", "assistant": "assistant"}
)
def formatting_prompts_func(examples):
convos = examples["conversations"]
texts = [tokenizer.apply_chat_template(c, tokenize=False, add_generation_prompt=False) for c in convos]
return {"text": texts}
dataset = dataset.map(formatting_prompts_func, batched=True)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
dataset_text_field="text",
max_seq_length=max_seq_length,
dataset_num_proc=2,
packing=False,
args=TrainingArguments(
per_device_train_batch_size=2,
gradient_accumulation_steps=4,
warmup_steps=5,
max_steps=60,
learning_rate=2e-4,
fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),
logging_steps=1,
optim="adamw_8bit",
weight_decay=0.01,
lr_scheduler_type="linear",
seed=42,
output_dir="outputs",
),
)
print("Starting SFT Training on InfraMind Trajectories...")
trainer_stats = trainer.train()
print("Training Complete!")
# Save LoRA adapters
model.save_pretrained("lora_model")
tokenizer.save_pretrained("lora_model")
print("✓ Model saved to lora_model/")
else:
print("\n--- MOCK TRAINING MODE ---")
print("Simulating training run because Unsloth/GPU is not available locally.")
print(f"Dataset Size: {len(dataset)} episodes")
print("Epoch 1/3 | Loss: 1.4532")
print("Epoch 2/3 | Loss: 0.8921")
print("Epoch 3/3 | Loss: 0.4201")
print("Training complete. (Run this in Google Colab for real GPU execution).")
if __name__ == "__main__":
main()
|