File size: 3,016 Bytes
67045e9 | 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 | #!/usr/bin/env python3
"""
LoRA fine-tuning for editorai:jtafull using Unsloth + TRL SFTTrainer.
Base model: Qwen/Qwen3-8B (fallback: Qwen/Qwen2.5-7B-Instruct)
"""
import json
import torch
from datasets import Dataset
from trl import SFTTrainer, SFTConfig
from unsloth import FastLanguageModel
from unsloth.chat_templates import get_chat_template
BASE_MODEL = "Qwen/Qwen3-8B"
MAX_SEQ_LENGTH = 4096
LORA_R = 16
LORA_ALPHA = 32
LORA_TARGETS = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"]
TRAIN_DATA = "data/train.jsonl"
OUTPUT_ADAPTER = "./lora-adapter"
OUTPUT_MERGED = "./merged-model"
def load_dataset(path: str) -> Dataset:
examples = []
with open(path) as f:
for line in f:
line = line.strip()
if line:
examples.append(json.loads(line))
print(f"Loaded {len(examples)} training examples")
return Dataset.from_list(examples)
def format_example(example, tokenizer):
return tokenizer.apply_chat_template(
example["messages"],
tokenize=False,
add_generation_prompt=False,
)
def main():
print(f"Loading base model: {BASE_MODEL}")
model, tokenizer = FastLanguageModel.from_pretrained(
model_name=BASE_MODEL,
max_seq_length=MAX_SEQ_LENGTH,
dtype=None,
load_in_4bit=True,
)
tokenizer = get_chat_template(tokenizer, chat_template="chatml")
model = FastLanguageModel.get_peft_model(
model,
r=LORA_R,
lora_alpha=LORA_ALPHA,
target_modules=LORA_TARGETS,
lora_dropout=0.05,
bias="none",
use_gradient_checkpointing="unsloth",
random_state=42,
)
dataset = load_dataset(TRAIN_DATA)
dataset = dataset.map(
lambda ex: {"text": format_example(ex, tokenizer)},
remove_columns=dataset.column_names,
)
trainer = SFTTrainer(
model=model,
tokenizer=tokenizer,
train_dataset=dataset,
args=SFTConfig(
dataset_text_field="text",
max_seq_length=MAX_SEQ_LENGTH,
per_device_train_batch_size=4,
gradient_accumulation_steps=4,
num_train_epochs=3,
learning_rate=2e-4,
lr_scheduler_type="cosine",
warmup_ratio=0.05,
fp16=not torch.cuda.is_bf16_supported(),
bf16=torch.cuda.is_bf16_supported(),
logging_steps=10,
save_steps=100,
output_dir=OUTPUT_ADAPTER,
report_to="none",
),
)
print("Starting training...")
trainer.train()
print(f"Saving LoRA adapter to {OUTPUT_ADAPTER}")
model.save_pretrained(OUTPUT_ADAPTER)
tokenizer.save_pretrained(OUTPUT_ADAPTER)
print(f"Merging LoRA into full model → {OUTPUT_MERGED}")
model.save_pretrained_merged(OUTPUT_MERGED, tokenizer, save_method="merged_16bit")
print("Done. Run convert_and_push.sh to build GGUF and push to Ollama.")
if __name__ == "__main__":
main()
|