Fix chat with LoRA

litgpt/chat/base.py

@@ -131,15 +131,15 @@ def main(
 
     fabric = L.Fabric(devices=1, precision=precision, plugins=plugins)
 
-    check_valid_checkpoint_dir(checkpoint_dir)
-    config = Config.from_file(checkpoint_dir / "model_config.yaml")
-
     checkpoint_path = checkpoint_dir / "lit_model.pth"
 
     # Merge if this is a raw LoRA checkpoint
-    if (checkpoint_path / "lit_model.pth.lora").is_file() and not checkpoint_path.is_file():
+    if (checkpoint_dir / "lit_model.pth.lora").is_file() and not checkpoint_path.is_file():
         print("Merging LoRA weights with the base model. This won't take long and is a one-time-only thing.")
-        merge_lora(checkpoint_path)
+        merge_lora(checkpoint_dir)
+
+    check_valid_checkpoint_dir(checkpoint_dir)
+    config = Config.from_file(checkpoint_dir / "model_config.yaml")
 
     with fabric.init_module(empty_init=True):
         model = GPT(config)

tests/test_chat.py

@@ -1,4 +1,5 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
+import os
 import re
 import subprocess
 import sys
@@ -129,3 +130,35 @@ def test_cli(mode):
     output = subprocess.check_output(args)
     output = str(output.decode())
     assert "Starts a conversation" in output
+
+
+@patch("litgpt.chat.base.input")
+@patch("litgpt.chat.base.merge_lora")
+def test_merge_lora_if_needed(mocked_merge_lora, mocked_input, fake_checkpoint_dir, monkeypatch, tensor_like):
+    # these values will be iteratively provided for each `input()` call
+    mocked_input.side_effect = [""]
+
+    # pretend there is an unmerged LORA checkpoint
+    os.rename(fake_checkpoint_dir / "lit_model.pth", fake_checkpoint_dir / "lit_model.pth.lora")
+    mocked_merge_lora.side_effect = lambda _: Path(fake_checkpoint_dir / "lit_model.pth").touch()
+
+    config_path = fake_checkpoint_dir / "model_config.yaml"
+    config = {
+        "name": "Llama 3",
+        "block_size": 128,
+        "vocab_size": 50,
+        "n_layer": 2,
+        "n_head": 4,
+        "n_embd": 8,
+        "rotary_percentage": 1,
+    }
+    config_path.write_text(yaml.dump(config))
+    monkeypatch.setattr(chat, "load_checkpoint", Mock())
+    monkeypatch.setattr(chat, "Tokenizer", Mock())
+
+    out, err = StringIO(), StringIO()
+    with redirect_stdout(out), redirect_stderr(err):
+        chat.main(checkpoint_dir=fake_checkpoint_dir)
+
+    assert re.match("Merging LoRA weights with the base model.", out.getvalue(), re.DOTALL)
+    mocked_merge_lora.assert_called_once()