Spaces:
Running
Running
| """ | |
| 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() | |