Inference and merging

Author: adaamko

scripts/merge_lora.py

@@ -1,85 +1,74 @@
-"""Merge a PEFT/LoRA checkpoint into its base causal LM.
+"""Merge a LoRA checkpoint into a standalone model.
 
-Example:
+Uses Unsloth for merging (same as training) to handle its internal patches
+correctly, and reproduces the exact tokenizer setup from training.
+
+Usage:
+    python scripts/merge_lora.py \
+        --checkpoint output/squeez_qwen/checkpoint-500 \
+        --output output/squeez_qwen_merged
+
+    # With explicit base model (if not auto-detected from adapter_config):
     python scripts/merge_lora.py \
-        --base-model Qwen/Qwen3.5-2B \
-        --adapter-path output/squeez_qwen/checkpoint-800 \
-        --output-dir output/squeez_qwen_merged
+        --checkpoint output/squeez_qwen \
+        --output output/squeez_qwen_merged \
+        --base-model Qwen/Qwen3.5-2B
 """
 
 from __future__ import annotations
 
 import argparse
+import json
 import logging
+from pathlib import Path
 
 logger = logging.getLogger(__name__)
 
 
-def build_parser() -> argparse.ArgumentParser:
-    parser = argparse.ArgumentParser(description="Merge a LoRA adapter into its base model")
-    parser.add_argument("--base-model", required=True, help="Base model name or path")
-    parser.add_argument("--adapter-path", required=True, help="Path to LoRA checkpoint")
-    parser.add_argument("--output-dir", required=True, help="Directory to save merged model")
-    parser.add_argument(
-        "--dtype",
-        choices=["auto", "bf16", "fp16", "fp32"],
-        default="auto",
-        help="Torch dtype to load the base model with before merging",
-    )
-    return parser
-
-
-def _resolve_dtype(dtype_name: str):
-    import torch
-
-    if dtype_name == "bf16":
-        return torch.bfloat16
-    if dtype_name == "fp16":
-        return torch.float16
-    if dtype_name == "fp32":
-        return torch.float32
-    if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
-        return torch.bfloat16
-    if torch.cuda.is_available():
-        return torch.float16
-    return torch.float32
-
-
 def main(argv: list[str] | None = None) -> int:
-    args = build_parser().parse_args(argv)
-
-    logging.basicConfig(
-        level=logging.INFO,
-        format="%(asctime)s [%(levelname)s] %(name)s: %(message)s",
+    parser = argparse.ArgumentParser(description="Merge LoRA checkpoint into standalone model")
+    parser.add_argument("--checkpoint", required=True, help="Path to LoRA checkpoint")
+    parser.add_argument("--output", required=True, help="Output path for merged model")
+    parser.add_argument("--base-model", default=None, help="Base model (auto-detected if omitted)")
+    parser.add_argument("--config", default=None, help="YAML config file")
+    args = parser.parse_args(argv)
+
+    logging.basicConfig(level=logging.INFO, format="%(message)s")
+
+    from unsloth import FastLanguageModel
+
+    from squeez.training.train import _prepare_text_tokenizer, load_config
+
+    config = load_config(args.config)
+    max_length = config.get("max_length", 16384)
+
+    # Detect base model from adapter config if not provided
+    base_model_name = args.base_model
+    adapter_config_path = Path(args.checkpoint) / "adapter_config.json"
+    if not base_model_name and adapter_config_path.exists():
+        with open(adapter_config_path) as f:
+            base_model_name = json.load(f).get("base_model_name_or_path", "")
+    if not base_model_name:
+        base_model_name = config.get("model", "Qwen/Qwen3.5-2B")
+
+    logger.info(f"Loading checkpoint from {args.checkpoint} (base: {base_model_name})")
+    model, tokenizer = FastLanguageModel.from_pretrained(
+        args.checkpoint,
+        max_seq_length=max_length,
+        load_in_4bit=False,
+        load_in_16bit=True,
     )
 
-    from peft import PeftModel
-    from transformers import AutoModelForCausalLM, AutoTokenizer
-
-    dtype = _resolve_dtype(args.dtype)
-
-    logger.info("Loading tokenizer from %s", args.base_model)
-    tokenizer = AutoTokenizer.from_pretrained(args.base_model, trust_remote_code=True)
+    # Reproduce the same tokenizer patches as training
+    tokenizer = _prepare_text_tokenizer(base_model_name, tokenizer)
 
-    logger.info("Loading base model from %s", args.base_model)
-    model = AutoModelForCausalLM.from_pretrained(
-        args.base_model,
-        torch_dtype=dtype,
-        trust_remote_code=True,
-        device_map="auto",
+    logger.info(f"Merging and saving to {args.output}")
+    model.save_pretrained_merged(
+        args.output,
+        tokenizer,
+        save_method="merged_16bit",
     )
-
-    logger.info("Loading adapter from %s", args.adapter_path)
-    model = PeftModel.from_pretrained(model, args.adapter_path)
-
-    logger.info("Merging adapter into base model")
-    model = model.merge_and_unload()
-
-    logger.info("Saving merged model to %s", args.output_dir)
-    model.save_pretrained(args.output_dir, safe_serialization=True)
-    tokenizer.save_pretrained(args.output_dir)
-
-    logger.info("Merge complete")
+    logger.info(f"Done. Merged model saved to {args.output}")
     return 0
 
 

squeez/inference/extractor.py

@@ -79,16 +79,6 @@ def _build_messages(task: str, tool_output: str) -> list[dict]:
     ]
 
 
-def _format_prompt(task: str, tool_output: str) -> str:
-    """Format the input prompt using the ChatML template for local generation."""
-    messages = _build_messages(task, tool_output)
-    return (
-        f"<|im_start|>system\n{messages[0]['content']}<|im_end|>\n"
-        f"<|im_start|>user\n{messages[1]['content']}<|im_end|>\n"
-        f"<|im_start|>assistant\n"
-    )
-
-
 def _is_encoder_model(model_path: str) -> bool:
     """Check if a model path contains a squeez-encoder model."""
     import json
@@ -385,7 +375,10 @@ def _extract_transformers(
         """Extract using local transformers model."""
         import torch
 
-        prompt = _format_prompt(task, tool_output)
+        messages = _build_messages(task, tool_output)
+        prompt = self._tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True
+        )
 
         inputs = self._tokenizer(
             prompt,

squeez/training/train.py

@@ -204,11 +204,12 @@ def train(args: argparse.Namespace):
     logger.info("Starting training...")
     trainer.train()
 
-    # 7. Save
-    logger.info(f"Saving model to {output_dir}")
-    trainer.save_model(output_dir)
+    # 7. Save merged model (LoRA weights folded into base — standalone, no adapter needed)
+    logger.info(f"Merging LoRA and saving full model to {output_dir}")
+    merged_model = model.merge_and_unload()
+    merged_model.save_pretrained(output_dir)
     tokenizer.save_pretrained(output_dir)
-    logger.info("Training complete!")
+    logger.info("Training complete! Saved merged model.")
 
 
 def build_parser(parser: argparse.ArgumentParser | None = None) -> argparse.ArgumentParser:

tests/test_extractor.py

@@ -1,33 +1,30 @@
 """Tests for squeez core functionality."""
 
 from squeez.data.config import SYSTEM_PROMPT
-from squeez.inference.extractor import _format_prompt, _load_config
+from squeez.inference.extractor import _build_messages, _load_config
 
 
-def test_format_prompt_basic():
-    prompt = _format_prompt("Fix the bug", "class Foo:\n    pass")
-    assert "Fix the bug" in prompt
-    assert "class Foo:" in prompt
-    assert SYSTEM_PROMPT in prompt
-    assert "<|im_start|>system" in prompt
-    assert "<|im_start|>user" in prompt
-    assert "<|im_start|>assistant" in prompt
-    assert "<|im_end|>" in prompt
+def test_build_messages_basic():
+    messages = _build_messages("Fix the bug", "class Foo:\n    pass")
+    assert len(messages) == 2
+    assert messages[0]["role"] == "system"
+    assert messages[0]["content"] == SYSTEM_PROMPT
+    assert messages[1]["role"] == "user"
+    assert "Fix the bug" in messages[1]["content"]
+    assert "class Foo:" in messages[1]["content"]
 
 
-def test_format_prompt_truncates_long_task():
+def test_build_messages_truncates_long_task():
     long_task = "x" * 5000
-    prompt = _format_prompt(long_task, "output")
-    assert len(long_task) > 3000
-    assert "..." in prompt
-    task_section = prompt.split("<query>\n", 1)[1].split("\n</query>", 1)[0]
+    messages = _build_messages(long_task, "output")
+    task_section = messages[1]["content"].split("<query>\n", 1)[1].split("\n</query>", 1)[0]
     assert len(task_section) == 3003  # 3000 + "..."
 
 
-def test_format_prompt_empty_task():
-    prompt = _format_prompt("", "some output")
-    assert "<query>" not in prompt
-    assert "some output" in prompt
+def test_build_messages_empty_task():
+    messages = _build_messages("", "some output")
+    assert "<query>" not in messages[1]["content"]
+    assert "some output" in messages[1]["content"]
 
 
 def test_system_prompt_has_relevant_lines_format():