finetune_qlora.py

import json
import time
from datetime import datetime
import torch

device = "cuda" if torch.cuda.is_available() else "cpu"
print(f"Using device: {device}")

from datasets import Dataset
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForLanguageModeling,
    BitsAndBytesConfig,
    TrainerCallback,
)
from peft import (
    get_peft_model,
    LoraConfig,
    TaskType,
    prepare_model_for_kbit_training,
)

# =========================
# Config
# =========================
OUTPUT_DIR = input("Enter the path where you want to save the finetuned model: ").strip()
epoch_num  = int(input("Number of epochs: ").strip())

MAX_LENGTH = 2048

# =========================
# Training time window
# (solar energy schedule)
# =========================
TRAINING_WINDOW_ENABLED = True
TRAINING_WINDOW_START   = 7
TRAINING_WINDOW_END     = 20

# =========================
# Time window helpers
# =========================
def is_within_training_window():
    if not TRAINING_WINDOW_ENABLED:
        return True
    hour = datetime.now().hour
    return TRAINING_WINDOW_START <= hour < TRAINING_WINDOW_END


def wait_for_training_window():
    if is_within_training_window():
        return
    now = datetime.now()
    print(f"\n⏸  Outside training window ({now.strftime('%H:%M')}). "
          f"Pausing until {TRAINING_WINDOW_START:02d}:00 ...")
    while not is_within_training_window():
        time.sleep(60)
    print(f"▶  Training window open ({datetime.now().strftime('%H:%M')}). Resuming ...")


class TimeWindowCallback(TrainerCallback):
    def on_step_end(self, args, state, control, **kwargs):
        if TRAINING_WINDOW_ENABLED and not is_within_training_window():
            wait_for_training_window()

# =========================
# Dataset builders
# =========================
def build_dataset_from_txt(path):
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()
    samples = [t.strip() for t in text.split("\n\n") if t.strip()]
    return Dataset.from_dict({"text": samples})


def build_dataset_from_chat_json(path):
    """Full JSON array — [{messages: [{role, content}]}]"""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    samples = []
    for convo in data:
        lines = []
        for msg in convo["messages"]:
            role = msg["role"].capitalize()
            lines.append(f"{role}: {msg['content']}")
        samples.append("\n".join(lines))
    return Dataset.from_dict({"text": samples})


def build_dataset_from_chat_jsonl(path):
    """Line-delimited JSONL — OpenAssistant / messages+input style"""
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            convo = json.loads(line)
            lines = []
            for msg in convo["messages"]:
                role = msg["role"].capitalize()
                if "input" in msg and msg["input"].strip():
                    lines.append(f"{role} [Input: {msg['input']}]: {msg['content']}")
                else:
                    lines.append(f"{role}: {msg['content']}")
            samples.append("\n".join(lines))
    return Dataset.from_dict({"text": samples})


def build_dataset_from_sharegpt_jsonl(path):
    """
    ShareGPT-style JSONL:
    {"conversations": [{"from": "system", "value": "..."}, {"from": "human", "value": "..."}, {"from": "gpt", "value": "..."}]}
    """
    role_map = {
        "system":    "System",
        "human":     "User",
        "gpt":       "Assistant",
        "assistant": "Assistant",
    }
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            convo = json.loads(line)
            turns = convo.get("conversations", [])
            lines = []
            for msg in turns:
                role  = role_map.get(msg.get("from", "").lower(),
                                     msg.get("from", "Unknown").capitalize())
                value = msg.get("value", "").strip()
                if value:
                    lines.append(f"{role}: {value}")
            if lines:
                samples.append("\n".join(lines))
    return Dataset.from_dict({"text": samples})


def build_dataset_from_alpaca_json(path):
    """Alpaca-style JSON array: [{"instruction": "...", "input": "...", "output": "..."}]"""
    with open(path, "r", encoding="utf-8") as f:
        data = json.load(f)
    samples = []
    for item in data:
        instruction = item.get("instruction", "").strip()
        inp         = item.get("input", "").strip()
        output      = item.get("output", "").strip()
        if inp:
            text = f"User: {instruction}\nInput: {inp}\nAssistant: {output}"
        else:
            text = f"User: {instruction}\nAssistant: {output}"
        if text.strip():
            samples.append(text)
    return Dataset.from_dict({"text": samples})


def build_dataset_from_alpaca_jsonl(path):
    """Alpaca-style JSONL — same schema as above but line-delimited."""
    samples = []
    with open(path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if not line:
                continue
            item        = json.loads(line)
            instruction = item.get("instruction", "").strip()
            inp         = item.get("input", "").strip()
            output      = item.get("output", "").strip()
            if inp:
                text = f"User: {instruction}\nInput: {inp}\nAssistant: {output}"
            else:
                text = f"User: {instruction}\nAssistant: {output}"
            if text.strip():
                samples.append(text)
    return Dataset.from_dict({"text": samples})


# =========================
# Dataset menu helper
# =========================
def pick_dataset():
    print("\nDataset type:")
    print("1) Plain .txt")
    print("2) Chat JSON (full array, messages/role/content)")
    print("3) Chat JSONL (line-delimited, OpenAssistant-style)")
    print("4) ShareGPT JSONL (conversations/from/value)")
    print("5) Alpaca JSON  (instruction/input/output array)")
    print("6) Alpaca JSONL (instruction/input/output line-delimited)")
    choice = input("Choice: ").strip()

    if choice == "1":
        path = input("Path to .txt file: ").strip()
        return build_dataset_from_txt(path)
    elif choice == "2":
        path = input("Path to chat JSON: ").strip()
        return build_dataset_from_chat_json(path)
    elif choice == "3":
        path = input("Path to chat JSONL: ").strip()
        return build_dataset_from_chat_jsonl(path)
    elif choice == "4":
        path = input("Path to ShareGPT JSONL: ").strip()
        return build_dataset_from_sharegpt_jsonl(path)
    elif choice == "5":
        path = input("Path to Alpaca JSON: ").strip()
        return build_dataset_from_alpaca_json(path)
    elif choice == "6":
        path = input("Path to Alpaca JSONL: ").strip()
        return build_dataset_from_alpaca_jsonl(path)
    else:
        print("❌ Invalid choice")
        return None

# =========================
# Preprocess factory
# =========================
def preprocess_factory(tokenizer):
    def preprocess(example):
        # Try ShareGPT conversations format first
        conversations = example.get("conversations")
        if conversations:
            text = ""
            for turn in conversations:
                frm = turn.get("from", "")
                val = turn.get("value", "")
                if frm == "system":
                    text += f"<|system|>{val}</s>"
                elif frm == "human":
                    text += f"<|user|>{val}</s>"
                elif frm == "gpt":
                    text += f"<|assistant|>{val}</s>"
        else:
            # Fallback: plain text field
            text = example.get("text", "")

        enc = tokenizer(
            text,
            truncation=True,
            padding="max_length",
            max_length=MAX_LENGTH,
        )
        enc["labels"] = enc["input_ids"].copy()
        return enc
    return preprocess

# =========================
# QLoRA finetune
# =========================
def finetune():
    base_model = input("\nBase model (HuggingFace ID or local path, e.g. mistralai/Mistral-7B-v0.1): ").strip()

    dataset = pick_dataset()
    if dataset is None:
        return

    # ── QLoRA quantisation config ──────────────────────────────────────────
    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    print(f"\n📦 Loading {base_model} in 4-bit …")
    tokenizer = AutoTokenizer.from_pretrained(base_model, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    model = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=bnb_config,
        device_map="auto",
    )

    # ── Prepare for k-bit training ─────────────────────────────────────────
    model = prepare_model_for_kbit_training(model)

    # ── LoRA config ────────────────────────────────────────────────────────
    print("\nLoRA settings (press Enter to use defaults):")
    r_val       = input("  LoRA rank r         [default 16]: ").strip() or "16"
    alpha_val   = input("  LoRA alpha          [default 32]: ").strip() or "32"
    dropout_val = input("  LoRA dropout        [default 0.05]: ").strip() or "0.05"

    lora_config = LoraConfig(
        r=int(r_val),
        lora_alpha=int(alpha_val),
        lora_dropout=float(dropout_val),
        bias="none",
        task_type=TaskType.CAUSAL_LM,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj",
                        "gate_proj", "up_proj", "down_proj"],
    )

    model = get_peft_model(model, lora_config)
    model.print_trainable_parameters()

    # ── Tokenise dataset ───────────────────────────────────────────────────
    dataset = dataset.map(
        preprocess_factory(tokenizer),
        remove_columns=dataset.column_names,
    )

    collator = DataCollatorForLanguageModeling(tokenizer=tokenizer, mlm=False)

    # ── Training args ──────────────────────────────────────────────────────
    args = TrainingArguments(
        output_dir=OUTPUT_DIR,
        num_train_epochs=epoch_num,
        per_device_train_batch_size=2,
        gradient_accumulation_steps=4,
        gradient_checkpointing=True,
        optim="paged_adamw_8bit",
        bf16=True,
        learning_rate=2e-4,
        lr_scheduler_type="cosine",
        warmup_ratio=0.03,
        logging_steps=50,
        save_strategy="epoch",
        report_to="none",
    )

    wait_for_training_window()

    trainer = Trainer(
        model=model,
        args=args,
        train_dataset=dataset,
        data_collator=collator,
        callbacks=[TimeWindowCallback()],
    )

    print("\n🚀 Starting QLoRA finetuning …")
    trainer.train()

    model.save_pretrained(OUTPUT_DIR)
    tokenizer.save_pretrained(OUTPUT_DIR)

    print(f"\n✅ Finetuning complete — adapter saved to {OUTPUT_DIR}")
    print("   To merge weights later, load the base model + adapter with PEFT and call merge_and_unload().")

# =========================
# Merge adapter → full model
# =========================
def merge():
    base_model   = input("Base model used during finetuning: ").strip()
    adapter_path = input("Path to saved adapter (OUTPUT_DIR): ").strip()
    merge_out    = input("Where to save the merged model: ").strip()

    from peft import PeftModel

    print(f"\n📦 Loading {base_model} in float16 for merging …")
    tokenizer = AutoTokenizer.from_pretrained(adapter_path, use_fast=True)

    base = AutoModelForCausalLM.from_pretrained(
        base_model,
        torch_dtype=torch.float16,
        device_map="auto",
    )

    print("🔗 Applying adapter …")
    model = PeftModel.from_pretrained(base, adapter_path)

    print("🔀 Merging weights …")
    merged = model.merge_and_unload()

    print(f"💾 Saving merged model to {merge_out} …")
    merged.save_pretrained(merge_out)
    tokenizer.save_pretrained(merge_out)

    print(f"\n✅ Merged model saved to {merge_out}")
    print("   You can now load it with AutoModelForCausalLM.from_pretrained() directly — no PEFT needed.")

# =========================
# Run (inference)
# =========================
def run():
    from peft import PeftModel

    base_model   = input("Base model used during finetuning: ").strip()
    adapter_path = input("Path to saved adapter (OUTPUT_DIR): ").strip()

    print(f"\n📦 Loading {base_model} + adapter …")

    bnb_config = BitsAndBytesConfig(
        load_in_4bit=True,
        bnb_4bit_use_double_quant=True,
        bnb_4bit_quant_type="nf4",
        bnb_4bit_compute_dtype=torch.bfloat16,
    )

    tokenizer = AutoTokenizer.from_pretrained(adapter_path, use_fast=True)
    tokenizer.pad_token = tokenizer.eos_token

    base = AutoModelForCausalLM.from_pretrained(
        base_model,
        quantization_config=bnb_config,
        device_map="auto",
    )
    model = PeftModel.from_pretrained(base, adapter_path)
    model.eval()

    print("\nModel ready. Type 'exit' to quit.")

    while True:
        user_prompt = input("\nUser > ").strip()
        if user_prompt.lower() == "exit":
            break

        system = input("System prompt (optional, leave empty to skip): ").strip()
        if system:
            prompt = f"<|system|>{system}</s><|user|>{user_prompt}</s><|assistant|>"
        else:
            prompt = f"<|user|>{user_prompt}</s><|assistant|>"

        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
        with torch.no_grad():
            out = model.generate(
                **inputs,
                max_new_tokens=200,
                do_sample=True,
                temperature=0.8,
                top_p=0.95,
                repetition_penalty=1.1,
            )
        print("\n" + tokenizer.decode(out[0], skip_special_tokens=True))

# =========================
# Menu
# =========================
def main():
    print("""
=============================
   QLoRA FINETUNING SCRIPT
=============================
1) Finetune (QLoRA)
2) Run (inference)
3) Merge adapter into base model
4) Exit
""")
    c = input("Select: ").strip()
    if c == "1":
        finetune()
    elif c == "2":
        run()
    elif c == "3":
        merge()
    else:
        print("👋 Bye")

if __name__ == "__main__":
    main()