QLoRA stuff + datasets

qlora/finetune_mistral_bali_qlora.py

@@ -0,0 +1,311 @@
+#!/usr/bin/env python3
+"""
+QLoRA SFT fine-tune for Mistral-7B on chat-style JSONL:
+Each line: {"messages": [{"role":"system","content":...}, {"role":"user","content":...}, {"role":"assistant","content":...}, ...]}
+
+Produces a LoRA adapter you can merge or load at inference time.
+
+Example:
+python finetune_mistral_bali_qlora.py \
+  --model_id mistralai/Mistral-7B-Instruct-v0.2 \
+  --train_jsonl /path/to/bali_train.jsonl \
+  --output_dir ./mistral-bali-lora \
+  --max_seq_len 2048 \
+  --per_device_train_batch_size 1 \
+  --gradient_accumulation_steps 16 \
+  --learning_rate 2e-4 \
+  --num_train_epochs 2 \
+  --streaming true
+"""
+
+import argparse
+import json
+import os
+from typing import Any, Dict, List, Optional
+
+import torch
+from datasets import load_dataset
+from peft import LoraConfig, get_peft_model, prepare_model_for_kbit_training
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    BitsAndBytesConfig,
+    TrainingArguments,
+)
+from trl import SFTTrainer
+
+
+# -----------------------------
+# Data formatting
+# -----------------------------
+def normalize_messages(messages: List[Dict[str, str]]) -> List[Dict[str, str]]:
+    """
+    Ensures message roles/content are well-formed and in allowed roles.
+    """
+    allowed = {"system", "user", "assistant"}
+    out = []
+    for m in messages:
+        role = (m.get("role") or "").strip().lower()
+        content = m.get("content")
+        if role not in allowed or content is None:
+            continue
+        content = str(content)
+        out.append({"role": role, "content": content})
+    return out
+
+
+def messages_to_text(tokenizer: AutoTokenizer, example: Dict[str, Any]) -> str:
+    """
+    Converts {"messages":[...]} to a single training text using the model's chat template if available.
+    For Mistral Instruct models, tokenizer.apply_chat_template is typically present.
+    """
+    messages = normalize_messages(example.get("messages", []))
+    if not messages:
+        return ""
+
+    # Prefer tokenizer chat template when available.
+    if (
+        hasattr(tokenizer, "apply_chat_template")
+        and tokenizer.chat_template is not None
+    ):
+        # add_generation_prompt=False -> include the assistant content in the formatted text
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=False,
+        )
+
+    # Fallback formatting (less ideal than the native template):
+    # Keep it deterministic and simple.
+    parts = []
+    for m in messages:
+        r = m["role"]
+        c = m["content"].strip()
+        if r == "system":
+            # Override system message
+            system_message = """
+You are a specialized Balinese cultural travel expert. Your role is to provide accurate, culturally grounded, and practical guidance for travelers engaging with Balinese culture, including temples, ceremonies, etiquette, ritual calendars, dance, crafts, village life, sacred landscapes, and historical–spiritual context.
+
+Prioritize cultural meaning and lived practice over sightseeing. Explain why places, rituals, and customs matter, and how visitors should behave respectfully. Emphasize dress codes, offerings, bodily conduct, photography rules, gender and purity considerations, and community norms.
+
+Integrate timing and context where relevant, including ceremonial cycles (Pawukon/Wuku, full and new moons), festival periods, tides, agricultural rhythms, and temple schedules. Promote responsible tourism, community benefit, and environmental care, and discourage entry into restricted or sacred spaces.
+
+Go beyond generic tips by naming specific temples, villages, regions, ceremonies, deities, and regional variations. Include practical logistics (access, hours, customary donations, crowd patterns) when helpful, without speculation. If uncertain, state this briefly and suggest local confirmation.
+
+Structure responses clearly: a brief contextual introduction, followed by well-labeled sections or bullet points, and a short “Essentials” or “Respect Checklist” summary.
+
+Do not include chain-of-thought, hidden reasoning, or meta commentary. Provide only polished, user-facing guidance in a calm, authoritative, and respectful tone.
+            """
+            parts.append(f"<<SYS>>\n{system_message}\n<</SYS>>\n")
+        elif r == "user":
+            parts.append(f"[USER]\n{c}\n")
+        else:
+            parts.append(f"[ASSISTANT]\n{c}\n")
+    return "\n".join(parts).strip() + "\n"
+
+
+# -----------------------------
+# Main
+# -----------------------------
+def parse_args():
+    p = argparse.ArgumentParser()
+    p.add_argument(
+        "--model_id",
+        type=str,
+        required=True,
+        help="e.g. mistralai/Mistral-7B-Instruct-v0.2 (recommended) or base model",
+    )
+    p.add_argument(
+        "--train_jsonl",
+        type=str,
+        required=True,
+        help="Path to JSONL training file; each line has a 'messages' list.",
+    )
+    p.add_argument(
+        "--eval_jsonl",
+        type=str,
+        default=None,
+        help="Optional eval JSONL with same format.",
+    )
+    p.add_argument("--output_dir", type=str, required=True)
+
+    # Training hyperparameters
+    p.add_argument("--max_seq_len", type=int, default=2048)
+    p.add_argument("--per_device_train_batch_size", type=int, default=1)
+    p.add_argument("--per_device_eval_batch_size", type=int, default=1)
+    p.add_argument("--gradient_accumulation_steps", type=int, default=16)
+    p.add_argument("--learning_rate", type=float, default=2e-4)
+    p.add_argument("--weight_decay", type=float, default=0.0)
+    p.add_argument("--num_train_epochs", type=float, default=1.0)
+    p.add_argument("--warmup_ratio", type=float, default=0.03)
+    p.add_argument("--logging_steps", type=int, default=10)
+    p.add_argument("--save_steps", type=int, default=200)
+    p.add_argument("--eval_steps", type=int, default=200)
+    p.add_argument("--seed", type=int, default=42)
+
+    # Performance / memory
+    p.add_argument(
+        "--streaming",
+        type=str,
+        default="true",
+        help="true/false. Use streaming for very large JSONL.",
+    )
+    p.add_argument(
+        "--bf16",
+        type=str,
+        default="true",
+        help="true/false. Prefer bf16 if your GPU supports it.",
+    )
+    p.add_argument("--gradient_checkpointing", type=str, default="true")
+
+    # LoRA config
+    p.add_argument("--lora_r", type=int, default=16)
+    p.add_argument("--lora_alpha", type=int, default=32)
+    p.add_argument("--lora_dropout", type=float, default=0.05)
+    p.add_argument(
+        "--target_modules",
+        type=str,
+        default="q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj",
+        help="Comma-separated module names for Mistral-style architectures.",
+    )
+
+    # Optional: limit samples for quick smoke tests
+    p.add_argument("--max_train_samples", type=int, default=None)
+    p.add_argument("--max_eval_samples", type=int, default=None)
+
+    return p.parse_args()
+
+
+def str2bool(x: str) -> bool:
+    return str(x).strip().lower() in {"1", "true", "yes", "y", "t"}
+
+
+def main():
+    args = parse_args()
+    os.makedirs(args.output_dir, exist_ok=True)
+
+    streaming = str2bool(args.streaming)
+    use_bf16 = str2bool(args.bf16)
+    use_gc = str2bool(args.gradient_checkpointing)
+
+    # -----------------------------
+    # Tokenizer
+    # -----------------------------
+    tokenizer = AutoTokenizer.from_pretrained(args.model_id, use_fast=True)
+    if tokenizer.pad_token is None:
+        # Common for causal LMs
+        tokenizer.pad_token = tokenizer.eos_token
+
+    # -----------------------------
+    # Model (4-bit QLoRA)
+    # -----------------------------
+    bnb_config = BitsAndBytesConfig(
+        load_in_4bit=True,
+        bnb_4bit_compute_dtype=torch.bfloat16 if use_bf16 else torch.float16,
+        bnb_4bit_quant_type="nf4",
+        bnb_4bit_use_double_quant=True,
+    )
+
+    model = AutoModelForCausalLM.from_pretrained(
+        args.model_id,
+        quantization_config=bnb_config,
+        device_map="auto",
+        torch_dtype=torch.bfloat16 if use_bf16 else torch.float16,
+    )
+
+    model.config.use_cache = False  # important for training
+    if use_gc:
+        model.gradient_checkpointing_enable()
+
+    # Prepare for k-bit + LoRA
+    model = prepare_model_for_kbit_training(model)
+
+    target_modules = [m.strip() for m in args.target_modules.split(",") if m.strip()]
+    lora_config = LoraConfig(
+        r=args.lora_r,
+        lora_alpha=args.lora_alpha,
+        lora_dropout=args.lora_dropout,
+        bias="none",
+        task_type="CAUSAL_LM",
+        target_modules=target_modules,
+    )
+    model = get_peft_model(model, lora_config)
+    model.print_trainable_parameters()
+
+    # -----------------------------
+    # Dataset
+    # -----------------------------
+    data_files = {"train": args.train_jsonl}
+    if args.eval_jsonl:
+        data_files["eval"] = args.eval_jsonl
+
+    ds = load_dataset("json", data_files=data_files, streaming=streaming)
+
+    def format_fn(example: Dict[str, Any]) -> Dict[str, str]:
+        text = messages_to_text(tokenizer, example)
+        return {"text": text}
+
+    train_ds = ds["train"].map(format_fn)
+    eval_ds = ds["eval"].map(format_fn) if args.eval_jsonl else None
+
+    # Optional sample limits (works differently for streaming vs non-streaming)
+    if args.max_train_samples is not None:
+        if streaming:
+            train_ds = train_ds.take(args.max_train_samples)
+        else:
+            train_ds = train_ds.select(
+                range(min(args.max_train_samples, len(train_ds)))
+            )
+
+    if eval_ds is not None and args.max_eval_samples is not None:
+        if streaming:
+            eval_ds = eval_ds.take(args.max_eval_samples)
+        else:
+            eval_ds = eval_ds.select(range(min(args.max_eval_samples, len(eval_ds))))
+
+    # -----------------------------
+    # Training
+    # -----------------------------
+    training_args = TrainingArguments(
+        output_dir=args.output_dir,
+        per_device_train_batch_size=args.per_device_train_batch_size,
+        per_device_eval_batch_size=args.per_device_eval_batch_size,
+        gradient_accumulation_steps=args.gradient_accumulation_steps,
+        learning_rate=args.learning_rate,
+        weight_decay=args.weight_decay,
+        num_train_epochs=args.num_train_epochs,
+        warmup_ratio=args.warmup_ratio,
+        logging_steps=args.logging_steps,
+        save_steps=args.save_steps,
+        eval_strategy="steps" if eval_ds is not None else "no",
+        eval_steps=args.eval_steps if eval_ds is not None else None,
+        save_total_limit=3,
+        bf16=use_bf16,
+        fp16=not use_bf16,
+        optim="paged_adamw_8bit",  # good default for QLoRA
+        lr_scheduler_type="cosine",
+        seed=args.seed,
+        report_to="none",
+    )
+
+    trainer = SFTTrainer(
+        model=model,
+        processing_class=tokenizer,
+        args=training_args,
+        train_dataset=train_ds,
+        eval_dataset=eval_ds,
+        packing=True,  # packs multiple conversations per sequence for higher throughput
+    )
+
+    trainer.train()
+
+    # Save LoRA adapter + tokenizer
+    trainer.model.save_pretrained(args.output_dir)
+    tokenizer.save_pretrained(args.output_dir)
+
+    print(f"\nDone. Saved LoRA adapter to: {args.output_dir}")
+    print("Inference: load base model + peft adapter from this directory.")
+
+
+if __name__ == "__main__":
+    main()

qlora/run.sh

@@ -0,0 +1,10 @@
+python finetune_mistral_bali_qlora.py \
+  --model_id mistralai/Mistral-7B-Instruct-v0.2 \
+  --train_jsonl ../raft/bali_culture_raft_dataset.jsonl \
+  --output_dir ./mistral-bali-lora \
+  --max_seq_len 2048 \
+  --per_device_train_batch_size 1 \
+  --gradient_accumulation_steps 16 \
+  --learning_rate 2e-4 \
+  --num_train_epochs 2 \
+  --streaming true

raft/jsonl_remapper.py

@@ -0,0 +1,138 @@
+#!/usr/bin/env python3
+"""
+Rewrite chat-style JSONL into {"input": ..., "output": ...} JSONL for LLM tuning.
+
+Expected input line shape (example):
+{
+  "messages": [
+    {"role":"system","content":"..."},
+    {"role":"user","content":"..."},
+    {"role":"assistant","content":"..."}
+  ],
+  "meta": {...}   # optional
+}
+
+Output line shape:
+{"input": "<user text>", "output": "<assistant text>"}
+
+By default:
+- Ignores all non-user/assistant roles (e.g., system).
+- Emits one record per (user -> next assistant) pair in the conversation.
+- Drops all other fields (including meta) unless --keep-meta is set.
+
+Usage:
+  python rewrite_jsonl.py in.jsonl out.jsonl
+  cat in.jsonl | python rewrite_jsonl.py - - > out.jsonl
+  python rewrite_jsonl.py in.jsonl out.jsonl --only-last
+  python rewrite_jsonl.py in.jsonl out.jsonl --keep-meta
+"""
+
+import argparse
+import json
+import sys
+from typing import Any, Dict, List, Optional, Tuple
+
+
+def iter_user_assistant_pairs(messages: List[Dict[str, Any]]) -> List[Tuple[str, str]]:
+    """
+    Return list of (user_content, assistant_content) pairs.
+    Pairing rule: whenever a 'user' message is followed later by the next 'assistant'
+    message, emit a pair. Intermediate system/tool messages are ignored.
+    """
+    pairs: List[Tuple[str, str]] = []
+    pending_user: Optional[str] = None
+
+    for m in messages:
+        role = m.get("role")
+        content = m.get("content")
+
+        if role == "user":
+            # Start (or restart) a pending user turn
+            if isinstance(content, str) and content.strip():
+                pending_user = content
+            else:
+                pending_user = ""
+        elif role == "assistant":
+            if pending_user is not None:
+                assistant_text = content if isinstance(content, str) else ""
+                pairs.append((pending_user, assistant_text))
+                pending_user = None
+        else:
+            # ignore system/tool/developer/etc.
+            continue
+
+    return pairs
+
+
+def read_lines(path: str) -> List[str]:
+    if path == "-":
+        return sys.stdin.read().splitlines()
+    with open(path, "r", encoding="utf-8") as f:
+        return f.read().splitlines()
+
+
+def write_lines(path: str, lines: List[str]) -> None:
+    if path == "-":
+        sys.stdout.write("\n".join(lines) + ("\n" if lines else ""))
+        return
+    with open(path, "w", encoding="utf-8") as f:
+        f.write("\n".join(lines) + ("\n" if lines else ""))
+
+
+def main() -> int:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("infile", help="Input JSONL path, or '-' for stdin")
+    ap.add_argument("outfile", help="Output JSONL path, or '-' for stdout")
+    ap.add_argument(
+        "--only-last",
+        action="store_true",
+        help="Emit only the last (user -> assistant) pair per input line.",
+    )
+    ap.add_argument(
+        "--keep-meta",
+        action="store_true",
+        help="If input line has 'meta', copy it through to output records.",
+    )
+    args = ap.parse_args()
+
+    in_lines = read_lines(args.infile)
+    out_lines: List[str] = []
+
+    for idx, line in enumerate(in_lines, start=1):
+        line = line.strip()
+        if not line:
+            continue
+
+        try:
+            obj = json.loads(line)
+        except json.JSONDecodeError as e:
+            sys.stderr.write(f"[line {idx}] JSON decode error: {e}\n")
+            continue
+
+        messages = obj.get("messages")
+        if not isinstance(messages, list):
+            # Not in expected format; skip silently (or log if desired)
+            continue
+
+        pairs = iter_user_assistant_pairs(messages)
+        if not pairs:
+            continue
+
+        if args.only_last:
+            pairs = [pairs[-1]]
+
+        for user_text, assistant_text in pairs:
+            out_obj: Dict[str, Any] = {
+                "input": user_text,
+                "output": assistant_text,
+            }
+            if args.keep_meta and isinstance(obj.get("meta"), dict):
+                out_obj["meta"] = obj["meta"]
+            out_lines.append(json.dumps(out_obj, ensure_ascii=False))
+
+    write_lines(args.outfile, out_lines)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

raft/jsonl_remapper_2.py

@@ -0,0 +1,36 @@
+import argparse
+import json
+
+
+def rewrite_jsonl(input_path, output_path):
+    with open(input_path, "r", encoding="utf-8") as infile, open(
+        output_path, "w", encoding="utf-8"
+    ) as outfile:
+
+        for line_num, line in enumerate(infile, start=1):
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                record = json.loads(line)
+                user_text = record.get("input", "")
+                bot_text = record.get("output", "")
+
+                new_record = {"text": f"<user>: {user_text} <bot>: {bot_text}"}
+
+                outfile.write(json.dumps(new_record, ensure_ascii=False) + "\n")
+
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON on line {line_num}") from e
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Rewrite JSONL from {input, output} to {text: '<user>: ... <bot>: ...'} format"
+    )
+    parser.add_argument("--input", required=True, help="Path to input JSONL file")
+    parser.add_argument("--output", required=True, help="Path to output JSONL file")
+
+    args = parser.parse_args()
+    rewrite_jsonl(args.input, args.output)

raft/jsonl_remapper_3.py

@@ -0,0 +1,54 @@
+import argparse
+import json
+
+
+def rewrite_jsonl(input_path, output_path):
+    with open(input_path, "r", encoding="utf-8") as infile, open(
+        output_path, "w", encoding="utf-8"
+    ) as outfile:
+
+        for line_num, line in enumerate(infile, start=1):
+            line = line.strip()
+            if not line:
+                continue
+
+            try:
+                record = json.loads(line)
+                messages = record.get("messages", [])
+
+                user_parts = []
+                bot_parts = []
+
+                for msg in messages:
+                    role = msg.get("role")
+                    content = msg.get("content", "")
+
+                    if role == "user":
+                        user_parts.append(content)
+                    elif role == "assistant":
+                        bot_parts.append(content)
+
+                # Skip entries without both sides
+                if not user_parts or not bot_parts:
+                    continue
+
+                user_text = " ".join(user_parts)
+                bot_text = " ".join(bot_parts)
+
+                new_record = {"text": f"<user>: {user_text} <bot>: {bot_text}"}
+
+                outfile.write(json.dumps(new_record, ensure_ascii=False) + "\n")
+
+            except json.JSONDecodeError as e:
+                raise ValueError(f"Invalid JSON on line {line_num}") from e
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Rewrite messages-based JSONL to {text: '<user>: ... <bot>: ...'} format"
+    )
+    parser.add_argument("--input", required=True, help="Path to input JSONL file")
+    parser.add_argument("--output", required=True, help="Path to output JSONL file")
+
+    args = parser.parse_args()
+    rewrite_jsonl(args.input, args.output)