135 lines
5.0 KiB
Python
135 lines
5.0 KiB
Python
"""
|
|
Training script: Fine-tune Qwen2.5-0.5B-Instruct for game command recognition from noisy STT.
|
|
|
|
Based on:
|
|
- arxiv:2502.12923 approach (structured JSON output, synthetic command data)
|
|
- TRL v1.3.0 SFTTrainer with LoRA (PEFT)
|
|
- Target: llama.cpp GGUF compatible output (Qwen2 architecture)
|
|
|
|
Dataset: Marco333/game-commands-noisy-stt (conversational messages format)
|
|
Model: Qwen/Qwen2.5-0.5B-Instruct -> fine-tuned -> merge LoRA -> push to Hub
|
|
|
|
Usage:
|
|
pip install trl peft transformers datasets accelerate torch
|
|
python train.py
|
|
"""
|
|
|
|
import torch
|
|
from datasets import load_dataset
|
|
from trl import SFTTrainer, SFTConfig
|
|
from peft import LoraConfig
|
|
|
|
# =============================================================================
|
|
# Configuration - edit these for your use case
|
|
# =============================================================================
|
|
|
|
MODEL_NAME = "Qwen/Qwen2.5-0.5B-Instruct"
|
|
DATASET_NAME = "Marco333/game-commands-noisy-stt"
|
|
OUTPUT_DIR = "./qwen2.5-0.5b-game-commands"
|
|
HUB_MODEL_ID = "Marco333/qwen2.5-0.5b-game-commands-stt"
|
|
|
|
# Training hyperparameters
|
|
LEARNING_RATE = 2e-4 # Higher LR for LoRA (10x base)
|
|
NUM_EPOCHS = 5
|
|
BATCH_SIZE = 8
|
|
GRAD_ACCUM_STEPS = 2
|
|
MAX_LENGTH = 512 # Short sequences for command recognition
|
|
WARMUP_RATIO = 0.05
|
|
|
|
# LoRA configuration
|
|
LORA_R = 32
|
|
LORA_ALPHA = 16
|
|
LORA_DROPOUT = 0.05
|
|
|
|
|
|
def main():
|
|
print(f"Training {MODEL_NAME} on {DATASET_NAME}")
|
|
dataset = load_dataset(DATASET_NAME)
|
|
|
|
peft_config = LoraConfig(
|
|
r=LORA_R,
|
|
lora_alpha=LORA_ALPHA,
|
|
lora_dropout=LORA_DROPOUT,
|
|
bias="none",
|
|
task_type="CAUSAL_LM",
|
|
target_modules=["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"],
|
|
)
|
|
|
|
training_args = SFTConfig(
|
|
output_dir=OUTPUT_DIR,
|
|
num_train_epochs=NUM_EPOCHS,
|
|
per_device_train_batch_size=BATCH_SIZE,
|
|
per_device_eval_batch_size=BATCH_SIZE,
|
|
gradient_accumulation_steps=GRAD_ACCUM_STEPS,
|
|
learning_rate=LEARNING_RATE,
|
|
lr_scheduler_type="cosine",
|
|
warmup_ratio=WARMUP_RATIO,
|
|
bf16=True,
|
|
max_length=MAX_LENGTH,
|
|
packing=False,
|
|
logging_steps=10,
|
|
logging_first_step=True,
|
|
disable_tqdm=True,
|
|
logging_strategy="steps",
|
|
eval_strategy="epoch",
|
|
save_strategy="epoch",
|
|
save_total_limit=2,
|
|
load_best_model_at_end=True,
|
|
metric_for_best_model="eval_loss",
|
|
push_to_hub=True,
|
|
hub_model_id=HUB_MODEL_ID,
|
|
hub_strategy="end",
|
|
report_to="none",
|
|
assistant_only_loss=True,
|
|
seed=42,
|
|
dataloader_num_workers=2,
|
|
)
|
|
|
|
trainer = SFTTrainer(
|
|
model=MODEL_NAME,
|
|
args=training_args,
|
|
train_dataset=dataset["train"],
|
|
eval_dataset=dataset["test"],
|
|
peft_config=peft_config,
|
|
)
|
|
|
|
print("Starting training...")
|
|
train_result = trainer.train()
|
|
print(f"Done! Loss: {train_result.training_loss:.4f}")
|
|
trainer.save_model()
|
|
|
|
# Merge LoRA into base model for GGUF compatibility
|
|
print("Merging LoRA...")
|
|
from peft import AutoPeftModelForCausalLM
|
|
from transformers import AutoTokenizer
|
|
|
|
merged = AutoPeftModelForCausalLM.from_pretrained(OUTPUT_DIR, dtype=torch.bfloat16, device_map="auto")
|
|
merged = merged.merge_and_unload()
|
|
merged_dir = f"{OUTPUT_DIR}-merged"
|
|
merged.save_pretrained(merged_dir)
|
|
tok = AutoTokenizer.from_pretrained(MODEL_NAME)
|
|
tok.save_pretrained(merged_dir)
|
|
|
|
#print("Pushing to Hub...")
|
|
#merged.push_to_hub(HUB_MODEL_ID, commit_message="Merged model - LoRA into Qwen2.5-0.5B for game commands")
|
|
#tok.push_to_hub(HUB_MODEL_ID, commit_message="Add tokenizer")
|
|
#print(f"Model: https://huggingface.co/{HUB_MODEL_ID}")
|
|
|
|
# Quick inference test
|
|
print("\nInference test:")
|
|
from transformers import pipeline
|
|
pipe = pipeline("text-generation", model=merged_dir, dtype=torch.bfloat16, device_map="auto")
|
|
system = "You are a game command interpreter. The user's input is noisy speech-to-text from a microphone in a gaming environment. Your job is to identify the intended game command from the noisy input.\n\nValid commands: Attack, CastSpell, CloseDoor, Crouch, Defend, DropItem, Heal, Interact, Jump, OpenDoor, OpenInventory, OpenMap, PauseGame, PickUp, Reload, SaveGame, Sprint, Stop, SwitchWeapon, UseItem\n\nRespond with ONLY a JSON object: {\"command\": \"<CommandName>\"}\nIf the input doesn't match any command, respond: {\"command\": \"Unknown\"}"
|
|
|
|
for inp in ["opun da dur", "attak enami", "hel me", "pic it op", "hello how are you"]:
|
|
out = pipe([{"role": "system", "content": system}, {"role": "user", "content": inp}], max_new_tokens=50, do_sample=False)
|
|
print(f" '{inp}' -> {out[0]['generated_text'][-1]['content']}")
|
|
|
|
print("\n\nGGUF conversion (for llama.cpp):")
|
|
print(f" python llama.cpp/convert_hf_to_gguf.py {HUB_MODEL_ID} --outtype f16 --outfile model.gguf")
|
|
print(f" ./llama-quantize model.gguf model_q4.gguf Q4_K_M")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|