Move scripts inside the package (#1045)

scripts/convert_hf_checkpoint.py → litgpt/scripts/convert_hf_checkpoint.py

@@ -2,20 +2,14 @@
 
 import gc
 import json
-import sys
 from collections import defaultdict
-from dataclasses import asdict
 from functools import partial
 from pathlib import Path
 from typing import Dict, List, Optional, Tuple, Union
 
 import torch
 from lightning.fabric.utilities.load import _NotYetLoadedTensor as NotYetLoadedTensor
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
 from litgpt import Config
 from litgpt.utils import incremental_save, lazy_load, save_config
 

scripts/convert_lit_checkpoint.py → litgpt/scripts/convert_lit_checkpoint.py

@@ -1,21 +1,16 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import gc
-import sys
 from functools import partial
 from pathlib import Path
 from typing import Dict, Optional, Tuple, Union
 
 import torch
 from lightning.fabric.utilities.load import _NotYetLoadedTensor as NotYetLoadedTensor
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
 from litgpt import Config
+from litgpt.scripts.convert_hf_checkpoint import layer_template, load_param
 from litgpt.utils import CLI, incremental_save, lazy_load
-from scripts.convert_hf_checkpoint import layer_template, load_param
 
 
 def copy_weights_falcon(

scripts/convert_pretrained_checkpoint.py → litgpt/scripts/convert_pretrained_checkpoint.py

@@ -1,17 +1,9 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
-import json
-import shutil
-import sys
-from dataclasses import asdict
 from pathlib import Path
 
 import torch
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
 from litgpt.utils import CLI, incremental_save, copy_config_files
 
 

scripts/download.py → litgpt/scripts/download.py

@@ -1,20 +1,15 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 
 import os
-import sys
 from contextlib import contextmanager
 from pathlib import Path
 from typing import Optional, List, Tuple
 
 import torch
 from lightning_utilities.core.imports import RequirementCache
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
+from litgpt.scripts.convert_hf_checkpoint import convert_hf_checkpoint
 from litgpt.utils import CLI
-from scripts.convert_hf_checkpoint import convert_hf_checkpoint
 
 _SAFETENSORS_AVAILABLE = RequirementCache("safetensors")
 _HF_TRANSFER_AVAILABLE = RequirementCache("hf_transfer")

scripts/merge_lora.py → litgpt/scripts/merge_lora.py

@@ -2,21 +2,15 @@
 
 """This script merges the LoRA weights with the base model"""
 import os
-import sys
 from pathlib import Path
 from typing import Optional, Tuple, Dict, Any
 
-import yaml
-
 import lightning as L
 import torch
-
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
+import yaml
 
 from litgpt.lora import GPT, Config, lora_filter, merge_lora_weights
-from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, lazy_load
+from litgpt.utils import CLI, check_valid_checkpoint_dir, lazy_load
 
 
 def merge_lora(

litgpt/utils.py

@@ -72,7 +72,7 @@ def check_valid_checkpoint_dir(checkpoint_dir: Path) -> None:
     error_message = (
         f"--checkpoint_dir {str(checkpoint_dir.absolute())!r}{problem}."
         "\nFind download instructions at https://github.com/Lightning-AI/litgpt/blob/main/tutorials\n"
-        f"{extra}\nSee all download options by running:\n python scripts/download.py"
+        f"{extra}\nSee all download options by running:\n python litgpt/scripts/download.py"
     )
     print(error_message, file=sys.stderr)
     raise SystemExit(1)

notebooks/falcon-inference.ipynb

@@ -60,7 +60,7 @@
    "outputs": [],
    "source": [
     "# download the weights\n",
-    "!python scripts/download.py --repo_id tiiuae/falcon-7b"
+    "!python litgpt/scripts/download.py --repo_id tiiuae/falcon-7b"
    ]
   },
   {

tests/test_adapter_v2.py

@@ -186,7 +186,7 @@ def test_against_hf_mixtral():
     from transformers.models.mixtral import MixtralConfig, MixtralForCausalLM
 
     from litgpt.adapter_v2 import GPT, Config
-    from scripts.convert_hf_checkpoint import copy_weights_hf_llama
+    from litgpt.scripts.convert_hf_checkpoint import copy_weights_hf_llama
 
     device = torch.device("cpu")
     dtype = torch.float32

tests/test_convert_hf_checkpoint.py

@@ -8,7 +8,7 @@
 
 def test_llama2_70b_conversion():
     from litgpt import Config
-    from scripts.convert_hf_checkpoint import copy_weights_hf_llama
+    from litgpt.scripts.convert_hf_checkpoint import copy_weights_hf_llama
 
     shapes = {
         "model.embed_tokens.weight": (32000, 8192),
@@ -102,14 +102,14 @@ def test_llama2_70b_conversion():
 
 
 def test_convert_hf_checkpoint(tmp_path):
-    from scripts.convert_hf_checkpoint import convert_hf_checkpoint
+    from litgpt.scripts.convert_hf_checkpoint import convert_hf_checkpoint
 
     with pytest.raises(ValueError, match="to contain .bin"):
         convert_hf_checkpoint(checkpoint_dir=tmp_path, model_name="pythia-14m")
 
     bin_file = tmp_path / "foo.bin"
     bin_file.touch()
-    with mock.patch("scripts.convert_hf_checkpoint.lazy_load") as load:
+    with mock.patch("litgpt.scripts.convert_hf_checkpoint.lazy_load") as load:
         convert_hf_checkpoint(checkpoint_dir=tmp_path, model_name="pythia-14m")
     load.assert_called_with(bin_file)
 

tests/test_convert_lit_checkpoint.py

@@ -16,7 +16,7 @@
 
 def test_convert_lit_checkpoint(tmp_path):
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import convert_lit_checkpoint
+    from litgpt.scripts.convert_lit_checkpoint import convert_lit_checkpoint
 
     ours_config = Config.from_name("Llama-2-7b-hf", block_size=8, n_layer=2, n_embd=32, n_head=2, padding_multiple=128)
     ours_model = GPT(ours_config)
@@ -43,7 +43,7 @@ def test_against_falcon_40b():
     from transformers.models.falcon.modeling_falcon import FalconForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_falcon as copy_to_theirs
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_falcon as copy_to_theirs
 
     ours_config = Config.from_name("falcon-40b", n_layer=2, n_head=8, n_query_groups=4, n_embd=32)
     theirs_config = FalconConfig(
@@ -78,7 +78,7 @@ def test_against_original_gpt_neox():
     from transformers import GPTNeoXConfig, GPTNeoXForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_gpt_neox as copy_to_theirs
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_gpt_neox as copy_to_theirs
 
     ours_config = Config(block_size=64, vocab_size=100, n_layer=4, n_head=8, n_embd=16)
     assert ours_config.padded_vocab_size == 512
@@ -123,7 +123,7 @@ def test_against_hf_llama2(ours_kwargs):
     from transformers.models.llama.modeling_llama import LlamaForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_llama
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_llama
 
     ours_config = Config.from_name(
         padded_vocab_size=10000, n_layer=2, n_head=8, n_embd=32, intermediate_size=86, **ours_kwargs
@@ -161,7 +161,7 @@ def test_against_mixtral():
     from transformers.models.mixtral import MixtralConfig, MixtralForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_llama
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_llama
 
     ours_config = Config.from_name(
         "Mixtral-8x7B-Instruct-v0.1",
@@ -208,7 +208,7 @@ def test_against_original_open_llama_3b():
     from transformers.models.llama.modeling_llama import LlamaForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_llama
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_llama
 
     ours_config = Config.from_name("open_llama_3b", n_layer=2, n_head=8, n_embd=32, intermediate_size=86)
     T = 5
@@ -250,7 +250,7 @@ def test_against_hf_phi_1_5():
             urlretrieve(url=url, filename=file_path)
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_phi
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_phi
     from reference_models.configuration_phi import PhiConfig
     from reference_models.original_phi_1_5 import PhiForCausalLM
 
@@ -300,7 +300,7 @@ def test_against_hf_phi_2():
             urlretrieve(url=url, filename=file_path)
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_phi
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_phi
     from reference_models.configuration_phi import PhiConfig
     from reference_models.original_phi_2 import PhiForCausalLM
 
@@ -341,7 +341,7 @@ def test_against_original_stablelm_zephyr_3b():
     from transformers import AutoConfig, AutoModelForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_llama
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_llama
 
     T = 5
     ours_config = Config.from_name("stablelm-zephyr-3b", n_layer=2, n_head=16, n_embd=32, intermediate_size=86)
@@ -395,7 +395,7 @@ def test_against_original_gemma(model_name, device, dtype):
     from transformers.models.gemma.modeling_gemma import GemmaForCausalLM
 
     from litgpt import GPT, Config
-    from scripts.convert_lit_checkpoint import copy_weights_llama
+    from litgpt.scripts.convert_lit_checkpoint import copy_weights_llama
 
     torch.set_default_dtype(dtype)
 
@@ -436,7 +436,7 @@ def test_against_original_gemma(model_name, device, dtype):
 
 
 def test_check_conversion_supported_adapter():
-    from scripts.convert_lit_checkpoint import check_conversion_supported
+    from litgpt.scripts.convert_lit_checkpoint import check_conversion_supported
 
     lit_weights = {"some.key.name": ANY, "error.key.gating_factor": ANY}
     with pytest.raises(NotImplementedError, match="Converting adapter"):
@@ -448,7 +448,7 @@ def test_check_conversion_supported_adapter():
 
 
 def test_check_conversion_supported_lora():
-    from scripts.convert_lit_checkpoint import check_conversion_supported
+    from litgpt.scripts.convert_lit_checkpoint import check_conversion_supported
 
     lit_weights = {"some.key.name": ANY, "error.key.lora": ANY}
     with pytest.raises(ValueError, match=r"LoRA.*cannot be converted"):
@@ -457,7 +457,7 @@ def test_check_conversion_supported_lora():
 
 def test_qkv_split():
     from litgpt import Config
-    from scripts.convert_lit_checkpoint import qkv_split
+    from litgpt.scripts.convert_lit_checkpoint import qkv_split
 
     # MHA
     config = Config(n_embd=4, n_head=4)

tests/test_convert_pretrained_checkpoint.py

@@ -6,7 +6,7 @@
 
 
 def test_convert_pretrained_checkpoint(tmp_path, fake_checkpoint_dir):
-    from scripts.convert_pretrained_checkpoint import convert_checkpoint
+    from litgpt.scripts.convert_pretrained_checkpoint import convert_checkpoint
 
     # Pretend we made a checkpoint from pretraining
     pretrained_checkpoint = {

tests/test_generate_sequentially.py

@@ -275,7 +275,7 @@ def test_model_forward_hooks():
 @RunIf(min_cuda_gpus=2)
 def test_base_with_sequentially(tmp_path):
     from litgpt import GPT, Config
-    from scripts.download import download_from_hub
+    from litgpt.scripts.download import download_from_hub
 
     # download the tokenizer
     download_from_hub(repo_id="EleutherAI/pythia-14m", tokenizer_only=True, checkpoint_dir=tmp_path)

tests/test_generate_tp.py

@@ -109,7 +109,7 @@ def test_tensor_parallel_llama(name, expected):
 @RunIf(min_cuda_gpus=2)
 def test_tp(tmp_path):
     from litgpt import GPT, Config
-    from scripts.download import download_from_hub
+    from litgpt.scripts.download import download_from_hub
 
     # download the tokenizer
     download_from_hub(repo_id="EleutherAI/pythia-14m", tokenizer_only=True, checkpoint_dir=tmp_path)

tests/test_lm_eval_harness.py

@@ -20,7 +20,7 @@ def test_run_eval(tmp_path, float_like):
     from eval.lm_eval_harness import EvalHarnessBase
     from litgpt.model import GPT
     from litgpt.tokenizer import Tokenizer
-    from scripts.download import download_from_hub
+    from litgpt.scripts.download import download_from_hub
 
     fabric = Fabric(devices=1)
     with fabric.init_module():

tests/test_lora.py

@@ -545,7 +545,7 @@ def test_against_hf_mixtral():
     from transformers.models.mixtral import MixtralConfig, MixtralForCausalLM
 
     from litgpt.lora import GPT, Config
-    from scripts.convert_hf_checkpoint import copy_weights_hf_llama
+    from litgpt.scripts.convert_hf_checkpoint import copy_weights_hf_llama
 
     device = torch.device("cpu")
     dtype = torch.float32

tests/test_merge_lora.py

@@ -20,7 +20,7 @@ def test_merge_lora(tmp_path, fake_checkpoint_dir):
     from litgpt.lora import GPT as LoRAGPT
     from litgpt.lora import lora_filter
     from litgpt.model import GPT
-    from scripts.merge_lora import merge_lora
+    from litgpt.scripts.merge_lora import merge_lora
 
     pretrained_checkpoint_dir = tmp_path / "pretrained"
     lora_checkpoint_dir = tmp_path / "lora"
@@ -74,7 +74,7 @@ def test_merge_lora(tmp_path, fake_checkpoint_dir):
 
 
 def test_load_lora_metadata(fake_checkpoint_dir):
-    from scripts.merge_lora import load_lora_metadata
+    from litgpt.scripts.merge_lora import load_lora_metadata
 
     assert not (fake_checkpoint_dir / "hyperparameters.yaml").is_file()
     with pytest.raises(FileNotFoundError, match="missing a `hyperparameters.yaml` file"):