Update ISO installer tiers and backend deployment
This commit is contained in:
Jino Jose committed
1 parent
79784a6743
commit
56668f7bdc
38 files changed
+7206
-175
No files matched your search
@@ -0,0 +1,309 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Nexus One AI — QLoRA Fine-Tuning Runner
|
||||
Launched as a subprocess by the FastAPI backend.
|
||||
|
||||
Writes structured JSONL log lines to --log-path so the UI can stream
|
||||
live loss curves and progress. Updates training_jobs.status in SQLite.
|
||||
|
||||
Requires (install on the training node):
|
||||
pip install torch transformers datasets peft bitsandbytes trl
|
||||
|
||||
Optional (faster, lower VRAM):
|
||||
pip install unsloth
|
||||
|
||||
Usage (called by main.py — do not run manually in production):
|
||||
python3 train_qlora.py --job-id 1 --db-path /opt/cezen/data/cezen.db \
|
||||
--dataset /opt/cezen/data/datasets/abc.jsonl \
|
||||
--base-model mistral:7b --output-dir /opt/cezen/data/finetuned/mymodel \
|
||||
--log-path /opt/cezen/data/job_logs/abc.jsonl \
|
||||
--epochs 3 --lr 2e-4 --batch-size 4 --lora-r 16 --lora-alpha 32 \
|
||||
--output-name mymodel
|
||||
"""
|
||||
|
||||
import argparse, json, os, sqlite3, sys, time
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
|
||||
# ── Argument parsing ──────────────────────────────────────────────────────────
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("--job-id", type=int, required=True)
|
||||
parser.add_argument("--db-path", required=True)
|
||||
parser.add_argument("--dataset", required=True)
|
||||
parser.add_argument("--base-model", required=True)
|
||||
parser.add_argument("--output-dir", required=True)
|
||||
parser.add_argument("--log-path", required=True)
|
||||
parser.add_argument("--output-name", required=True)
|
||||
parser.add_argument("--epochs", type=int, default=3)
|
||||
parser.add_argument("--lr", type=float, default=2e-4)
|
||||
parser.add_argument("--batch-size", type=int, default=4)
|
||||
parser.add_argument("--lora-r", type=int, default=16)
|
||||
parser.add_argument("--lora-alpha", type=int, default=32)
|
||||
args = parser.parse_args()
|
||||
|
||||
# ── Helpers ───────────────────────────────────────────────────────────────────
|
||||
|
||||
def utcnow():
|
||||
return datetime.now(timezone.utc).isoformat()
|
||||
|
||||
def db_connect():
|
||||
conn = sqlite3.connect(args.db_path)
|
||||
conn.row_factory = sqlite3.Row
|
||||
return conn
|
||||
|
||||
log_file = open(args.log_path, "a", buffering=1)
|
||||
|
||||
def log(type_: str, **kwargs):
|
||||
entry = {"ts": utcnow(), "type": type_, **kwargs}
|
||||
log_file.write(json.dumps(entry) + "\n")
|
||||
|
||||
def set_status(status: str):
|
||||
db = db_connect()
|
||||
if status in ("completed", "failed", "cancelled"):
|
||||
db.execute(
|
||||
"UPDATE training_jobs SET status=?, finished_at=? WHERE id=?",
|
||||
(status, utcnow(), args.job_id)
|
||||
)
|
||||
else:
|
||||
db.execute("UPDATE training_jobs SET status=? WHERE id=?", (status, args.job_id))
|
||||
db.commit()
|
||||
db.close()
|
||||
|
||||
# ── Dataset loading ───────────────────────────────────────────────────────────
|
||||
|
||||
def load_dataset_from_file(path: str):
|
||||
"""Load JSONL or CSV dataset into a list of dicts with 'text' or 'prompt'/'completion' keys."""
|
||||
p = Path(path)
|
||||
rows = []
|
||||
if p.suffix.lower() == ".csv":
|
||||
import csv
|
||||
with open(path, newline="", encoding="utf-8", errors="replace") as f:
|
||||
reader = csv.DictReader(f)
|
||||
for row in reader:
|
||||
rows.append(dict(row))
|
||||
else:
|
||||
with open(path, encoding="utf-8", errors="replace") as f:
|
||||
for line in f:
|
||||
line = line.strip()
|
||||
if line:
|
||||
try:
|
||||
rows.append(json.loads(line))
|
||||
except Exception:
|
||||
pass
|
||||
return rows
|
||||
|
||||
def format_row(row: dict) -> str:
|
||||
"""Convert a dataset row to a plain text training string."""
|
||||
if "text" in row:
|
||||
return row["text"]
|
||||
if "prompt" in row and "completion" in row:
|
||||
return f"### Instruction:\n{row['prompt']}\n\n### Response:\n{row['completion']}"
|
||||
if "instruction" in row and "output" in row:
|
||||
inp = row.get("input", "")
|
||||
return (f"### Instruction:\n{row['instruction']}\n\n### Input:\n{inp}\n\n### Response:\n{row['output']}"
|
||||
if inp else
|
||||
f"### Instruction:\n{row['instruction']}\n\n### Response:\n{row['output']}")
|
||||
# Fallback: concatenate all values
|
||||
return " ".join(str(v) for v in row.values() if v)
|
||||
|
||||
# ── Main training routine ─────────────────────────────────────────────────────
|
||||
|
||||
def main():
|
||||
log("start", job_id=args.job_id, base_model=args.base_model,
|
||||
epochs=args.epochs, lr=args.lr, batch_size=args.batch_size,
|
||||
lora_r=args.lora_r, lora_alpha=args.lora_alpha)
|
||||
set_status("running")
|
||||
|
||||
# Resolve model name (Ollama uses "mistral:7b" style — strip the tag for HF)
|
||||
hf_model = args.base_model
|
||||
if ":" in hf_model and "/" not in hf_model:
|
||||
# e.g. "mistral:7b" → try to map to HF repo
|
||||
name_map = {
|
||||
"mistral": "mistralai/Mistral-7B-v0.1",
|
||||
"llama2": "meta-llama/Llama-2-7b-hf",
|
||||
"llama3": "meta-llama/Meta-Llama-3-8B",
|
||||
"phi3": "microsoft/Phi-3-mini-4k-instruct",
|
||||
"gemma": "google/gemma-7b",
|
||||
"codellama":"codellama/CodeLlama-7b-hf",
|
||||
"qwen2": "Qwen/Qwen2-7B",
|
||||
}
|
||||
base_name = hf_model.split(":")[0].lower()
|
||||
hf_model = name_map.get(base_name, hf_model)
|
||||
log("info", msg=f"Mapped '{args.base_model}' → '{hf_model}' (HuggingFace)")
|
||||
|
||||
# Load dataset
|
||||
log("info", msg="Loading dataset...")
|
||||
raw_rows = load_dataset_from_file(args.dataset)
|
||||
if not raw_rows:
|
||||
log("error", msg="Dataset is empty or could not be parsed")
|
||||
set_status("failed")
|
||||
sys.exit(1)
|
||||
|
||||
texts = [format_row(r) for r in raw_rows]
|
||||
log("info", msg=f"Loaded {len(texts)} training examples")
|
||||
|
||||
# Try Unsloth first (faster), fall back to HF PEFT
|
||||
use_unsloth = False
|
||||
try:
|
||||
from unsloth import FastLanguageModel
|
||||
use_unsloth = True
|
||||
log("info", msg="Using Unsloth for accelerated training")
|
||||
except ImportError:
|
||||
log("info", msg="Unsloth not available — using HuggingFace PEFT + BitsAndBytes")
|
||||
|
||||
try:
|
||||
import torch
|
||||
from transformers import TrainingArguments, TrainerCallback
|
||||
from datasets import Dataset as HFDataset
|
||||
|
||||
if use_unsloth:
|
||||
model, tokenizer = FastLanguageModel.from_pretrained(
|
||||
model_name=hf_model,
|
||||
max_seq_length=2048,
|
||||
dtype=None,
|
||||
load_in_4bit=True,
|
||||
)
|
||||
model = FastLanguageModel.get_peft_model(
|
||||
model,
|
||||
r=args.lora_r,
|
||||
lora_alpha=args.lora_alpha,
|
||||
target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
|
||||
lora_dropout=0,
|
||||
bias="none",
|
||||
use_gradient_checkpointing="unsloth",
|
||||
)
|
||||
else:
|
||||
from transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
|
||||
from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
|
||||
|
||||
bnb_cfg = BitsAndBytesConfig(
|
||||
load_in_4bit=True,
|
||||
bnb_4bit_quant_type="nf4",
|
||||
bnb_4bit_compute_dtype=torch.float16,
|
||||
)
|
||||
tokenizer = AutoTokenizer.from_pretrained(hf_model, trust_remote_code=True)
|
||||
if tokenizer.pad_token is None:
|
||||
tokenizer.pad_token = tokenizer.eos_token
|
||||
|
||||
model = AutoModelForCausalLM.from_pretrained(
|
||||
hf_model,
|
||||
quantization_config=bnb_cfg,
|
||||
device_map="auto",
|
||||
trust_remote_code=True,
|
||||
)
|
||||
model = prepare_model_for_kbit_training(model)
|
||||
|
||||
lora_cfg = LoraConfig(
|
||||
r=args.lora_r,
|
||||
lora_alpha=args.lora_alpha,
|
||||
target_modules=["q_proj","k_proj","v_proj","o_proj"],
|
||||
lora_dropout=0.05,
|
||||
bias="none",
|
||||
task_type="CAUSAL_LM",
|
||||
)
|
||||
model = get_peft_model(model, lora_cfg)
|
||||
|
||||
# Tokenise
|
||||
def tokenise(examples):
|
||||
return tokenizer(
|
||||
examples["text"],
|
||||
truncation=True,
|
||||
max_length=2048,
|
||||
padding="max_length",
|
||||
)
|
||||
|
||||
hf_ds = HFDataset.from_dict({"text": texts})
|
||||
hf_ds = hf_ds.map(tokenise, batched=True, remove_columns=["text"])
|
||||
|
||||
# Custom callback to stream loss to our log
|
||||
class LossLogger(TrainerCallback):
|
||||
def on_log(self, _args, state, control, logs=None, **kwargs):
|
||||
if logs and "loss" in logs:
|
||||
log("loss",
|
||||
step=state.global_step,
|
||||
loss=round(float(logs["loss"]), 6),
|
||||
epoch=round(float(logs.get("epoch", 0)), 3),
|
||||
lr=float(logs.get("learning_rate", args.lr)))
|
||||
|
||||
output_dir = args.output_dir
|
||||
Path(output_dir).mkdir(parents=True, exist_ok=True)
|
||||
|
||||
from trl import SFTTrainer
|
||||
trainer = SFTTrainer(
|
||||
model=model,
|
||||
tokenizer=tokenizer,
|
||||
train_dataset=hf_ds,
|
||||
dataset_text_field="input_ids",
|
||||
max_seq_length=2048,
|
||||
args=TrainingArguments(
|
||||
output_dir=output_dir,
|
||||
num_train_epochs=args.epochs,
|
||||
per_device_train_batch_size=args.batch_size,
|
||||
gradient_accumulation_steps=4,
|
||||
warmup_steps=5,
|
||||
learning_rate=args.lr,
|
||||
fp16=not torch.cuda.is_bf16_supported(),
|
||||
bf16=torch.cuda.is_bf16_supported(),
|
||||
logging_steps=1,
|
||||
save_strategy="epoch",
|
||||
report_to="none",
|
||||
),
|
||||
callbacks=[LossLogger()],
|
||||
)
|
||||
|
||||
log("info", msg="Training started")
|
||||
trainer.train()
|
||||
log("info", msg="Training complete — saving model")
|
||||
trainer.save_model(output_dir)
|
||||
tokenizer.save_pretrained(output_dir)
|
||||
|
||||
except Exception as e:
|
||||
import traceback
|
||||
log("error", msg=str(e), traceback=traceback.format_exc())
|
||||
set_status("failed")
|
||||
sys.exit(1)
|
||||
|
||||
# Auto-register with Ollama via Modelfile
|
||||
try:
|
||||
_register_with_ollama(output_dir, args.output_name)
|
||||
except Exception as e:
|
||||
log("warning", msg=f"Could not auto-register with Ollama: {e}")
|
||||
|
||||
log("complete", msg="Job finished successfully", output_dir=output_dir)
|
||||
set_status("completed")
|
||||
|
||||
|
||||
def _register_with_ollama(model_dir: str, model_name: str):
|
||||
"""Create an Ollama Modelfile and register the fine-tuned model."""
|
||||
modelfile_path = Path(model_dir) / "Modelfile"
|
||||
modelfile_path.write_text(
|
||||
f'FROM {model_dir}\n'
|
||||
f'PARAMETER stop "<|im_end|>"\n'
|
||||
f'SYSTEM "This is a Nexus One AI fine-tuned model."\n'
|
||||
)
|
||||
import subprocess
|
||||
result = subprocess.run(
|
||||
["ollama", "create", model_name, "-f", str(modelfile_path)],
|
||||
capture_output=True, text=True, timeout=300
|
||||
)
|
||||
if result.returncode == 0:
|
||||
log("info", msg=f"Model '{model_name}' registered with Ollama")
|
||||
else:
|
||||
log("warning", msg=f"Ollama registration failed: {result.stderr}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
log("error", msg="Job interrupted (SIGTERM/SIGINT)")
|
||||
set_status("cancelled")
|
||||
sys.exit(130)
|
||||
except Exception as e:
|
||||
import traceback
|
||||
log("error", msg=str(e), traceback=traceback.format_exc())
|
||||
set_status("failed")
|
||||
sys.exit(1)
|
||||
finally:
|
||||
log_file.close()
|
||||
Reference in new issue
Block a user