Move generate/ inside the package (#1047)

README.md

@@ -108,7 +108,7 @@ To generate text predictions, you need to download the model weights. **If you d
 Run inference:
 
 ```bash
-python generate/base.py --prompt "Hello, my name is"
+python litgpt/generate/base.py --prompt "Hello, my name is"
 ```
 
 This will run the 3B pretrained model and require ~7 GB of GPU memory using the `bfloat16` datatype.

chat/base.py

@@ -13,7 +13,7 @@
 wd = Path(__file__).parent.parent.resolve()
 sys.path.append(str(wd))
 
-from generate.base import next_token
+from litgpt.generate.base import next_token
 from litgpt import GPT, Config, PromptStyle, Tokenizer
 from litgpt.prompts import load_prompt_style, has_prompt_style
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, load_checkpoint

eval/lm_eval_harness.py

@@ -15,7 +15,7 @@
 wd = Path(__file__).parent.parent.resolve()
 sys.path.append(str(wd))
 
-from generate.base import generate
+from litgpt.generate.base import generate
 from litgpt import GPT, Config, Tokenizer
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, load_checkpoint
 

litgpt/data/__init__.py

@@ -17,7 +17,7 @@
 
 __all__ = [
     "Alpaca",
-    "Alpaca2k"
+    "Alpaca2k",
     "AlpacaGPT4",
     "Deita",
     "Dolly",
@@ -30,6 +30,5 @@
     "SFTDataset",
     "TinyLlama",
     "TinyStories",
-    "apply_prompt_template",
     "get_sft_collate_fn",
 ]

litgpt/finetune/adapter.py

@@ -1,7 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import dataclasses
 import os
-import sys
 import time
 from pathlib import Path
 from pprint import pprint
@@ -18,6 +17,7 @@
 from litgpt.adapter import GPT, Block, Config, adapter_filter, mark_only_adapter_as_trainable
 from litgpt.args import EvalArgs, TrainArgs
 from litgpt.data import Alpaca, LitDataModule
+from litgpt.generate.base import generate
 from litgpt.prompts import save_prompt_style
 from litgpt.tokenizer import Tokenizer
 from litgpt.utils import (
@@ -33,12 +33,6 @@
     save_hyperparameters,
 )
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate.base import generate
-
 
 def setup(
     precision: Optional[str] = None,

litgpt/finetune/adapter_v2.py

@@ -1,7 +1,6 @@
 # Copyright Lightning AI. Licensed under the Apache License 2.0, see LICENSE file.
 import dataclasses
 import os
-import sys
 import time
 from pathlib import Path
 from pprint import pprint
@@ -18,6 +17,7 @@
 from litgpt.adapter_v2 import GPT, Block, Config, adapter_filter, mark_only_adapter_v2_as_trainable
 from litgpt.args import EvalArgs, TrainArgs
 from litgpt.data import Alpaca, LitDataModule
+from litgpt.generate.base import generate
 from litgpt.prompts import save_prompt_style
 from litgpt.tokenizer import Tokenizer
 from litgpt.utils import (
@@ -33,12 +33,6 @@
     save_hyperparameters,
 )
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate.base import generate
-
 
 def setup(
     precision: Optional[str] = None,

litgpt/finetune/full.py

@@ -2,7 +2,6 @@
 import dataclasses
 import math
 import os
-import sys
 import time
 from pathlib import Path
 from pprint import pprint
@@ -17,6 +16,7 @@
 
 from litgpt.args import EvalArgs, TrainArgs
 from litgpt.data import Alpaca, LitDataModule
+from litgpt.generate.base import generate
 from litgpt.model import GPT, Block, Config
 from litgpt.prompts import save_prompt_style
 from litgpt.tokenizer import Tokenizer
@@ -33,12 +33,6 @@
     save_hyperparameters,
 )
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate.base import generate
-
 
 def setup(
     precision: Optional[str] = None,

litgpt/finetune/lora.py

@@ -2,7 +2,6 @@
 import dataclasses
 import math
 import os
-import sys
 import time
 from pathlib import Path
 from pprint import pprint
@@ -19,6 +18,7 @@
 
 from litgpt.args import EvalArgs, TrainArgs
 from litgpt.data import LitDataModule, Alpaca
+from litgpt.generate.base import generate
 from litgpt.lora import GPT, Block, Config, lora_filter, mark_only_lora_as_trainable
 from litgpt.prompts import save_prompt_style
 from litgpt.tokenizer import Tokenizer
@@ -35,12 +35,6 @@
     save_hyperparameters,
 )
 
-# support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
-sys.path.append(str(wd))
-
-from generate.base import generate
-
 
 def setup(
     precision: Optional[str] = None,

generate/adapter.py → litgpt/generate/adapter.py

@@ -10,17 +10,16 @@
 from lightning.fabric.plugins import BitsandbytesPrecision
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-from generate.base import generate
 from litgpt import Tokenizer, PromptStyle
 from litgpt.adapter import GPT, Config
+from litgpt.generate.base import generate
 from litgpt.prompts import load_prompt_style, has_prompt_style
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, lazy_load
 
 
-
 def main(
     prompt: str = "What food do llamas eat?",
     input: str = "",

generate/adapter_v2.py → litgpt/generate/adapter_v2.py

@@ -10,12 +10,12 @@
 from lightning.fabric.plugins import BitsandbytesPrecision
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-from generate.base import generate
 from litgpt import Tokenizer, PromptStyle
 from litgpt.adapter_v2 import GPT, Config
+from litgpt.generate.base import generate
 from litgpt.prompts import load_prompt_style, has_prompt_style
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, lazy_load
 

generate/base.py → litgpt/generate/base.py

@@ -9,10 +9,9 @@
 import torch
 import torch._dynamo.config
 import torch._inductor.config
-from lightning.fabric.plugins import BitsandbytesPrecision
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
 from litgpt import GPT, Config, Tokenizer, PromptStyle

generate/full.py → litgpt/generate/full.py

@@ -10,11 +10,11 @@
 from lightning.fabric.plugins import BitsandbytesPrecision
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-from generate.base import generate
 from litgpt import GPT, Config, Tokenizer, PromptStyle
+from litgpt.generate.base import generate
 from litgpt.prompts import load_prompt_style, has_prompt_style
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, load_checkpoint
 

generate/lora.py → litgpt/generate/lora.py

@@ -10,11 +10,11 @@
 from lightning.fabric.plugins import BitsandbytesPrecision
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-from generate.base import generate
 from litgpt import Tokenizer, PromptStyle
+from litgpt.generate.base import generate
 from litgpt.lora import GPT, Config, merge_lora_weights
 from litgpt.prompts import load_prompt_style, has_prompt_style
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision, lazy_load

generate/sequentially.py → litgpt/generate/sequentially.py

@@ -18,10 +18,10 @@
 from typing_extensions import Type
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-import generate.base as generate_base
+import litgpt.generate.base as generate_base
 from litgpt import GPT, Config, Tokenizer
 from litgpt.model import Block, build_mask_cache
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision

generate/tp.py → litgpt/generate/tp.py

@@ -16,10 +16,10 @@
 from torch.distributed._functional_collectives import all_reduce
 
 # support running without installing as a package
-wd = Path(__file__).parent.parent.resolve()
+wd = Path(__file__).parents[2].resolve()
 sys.path.append(str(wd))
 
-import generate.base as generate_base
+import litgpt.generate.base as generate_base
 from litgpt import GPT, Config, Tokenizer
 from litgpt.model import CausalSelfAttention, GptNeoxMLP, LLaMAMLP, LLaMAMoE
 from litgpt.utils import CLI, check_valid_checkpoint_dir, get_default_supported_precision

notebooks/falcon-inference.ipynb

@@ -89,7 +89,7 @@
    ],
    "source": [
     "# run inference\n",
-    "!python generate/base.py \\\n",
+    "!python litgpt/generate/base.py \\\n",
     "        --prompt \"Hello, my name is\" \\\n",
     "        --checkpoint_dir checkpoints/tiiuae/falcon-7b \\\n",
     "        --quantize bnb.int8"

tests/test_chat.py

@@ -26,7 +26,7 @@
 )
 def test_generate(monkeypatch, generated, stop_tokens, expected):
     import chat.base as chat
-    import generate.base as generate
+    import litgpt.generate.base as generate
 
     input_idx = torch.tensor([5, 3])
     max_returned_tokens = len(input_idx) + 8

tests/test_generate.py

@@ -17,7 +17,7 @@
     "max_seq_length", (pytest.param(10, marks=pytest.mark.xfail(raises=NotImplementedError, strict=True)), 20 + 5)
 )
 def test_generate(max_seq_length):
-    import generate.base as generate
+    import litgpt.generate.base as generate
     from litgpt import GPT, Config
 
     T = 5
@@ -36,7 +36,7 @@ def multinomial(*args, **kwargs):
         multinomial_results.append(out)
         return out
 
-    with mock.patch("generate.base.multinomial_num_samples_1", multinomial):
+    with mock.patch("litgpt.generate.base.multinomial_num_samples_1", multinomial):
         out = generate.generate(model, input_idx, T + max_new_tokens, top_k=4)
 
     assert out.size(0) == T + max_new_tokens
@@ -47,7 +47,7 @@ def multinomial(*args, **kwargs):
 
 
 def test_main(fake_checkpoint_dir, monkeypatch, tensor_like):
-    import generate.base as generate
+    import litgpt.generate.base as generate
 
     config_path = fake_checkpoint_dir / "lit_config.json"
     config = {"block_size": 128, "vocab_size": 50, "n_layer": 2, "n_head": 4, "n_embd": 8, "rotary_percentage": 1}
@@ -85,15 +85,15 @@ def test_main(fake_checkpoint_dir, monkeypatch, tensor_like):
 
 
 def test_cli():
-    cli_path = Path(__file__).parent.parent / "generate" / "base.py"
+    cli_path = Path(__file__).parent.parent / "litgpt/generate/base.py"
     output = subprocess.check_output([sys.executable, cli_path, "-h"])
     output = str(output.decode())
     assert "Generates text samples" in output
 
 
 @pytest.mark.parametrize("temperature", (0.0, 1.0, 0.5))
 def test_sample(temperature):
-    from generate.base import sample
+    from litgpt.generate.base import sample
 
     # shape: 2x3x5
     logits = torch.tensor(

tests/test_generate_adapter.py

@@ -15,9 +15,9 @@
 @pytest.mark.parametrize("version", ("v1", "v2"))
 def test_main(fake_checkpoint_dir, monkeypatch, version, tensor_like):
     if version == "v1":
-        import generate.adapter as generate
+        import litgpt.generate.adapter as generate
     else:
-        import generate.adapter_v2 as generate
+        import litgpt.generate.adapter_v2 as generate
 
     config_path = fake_checkpoint_dir / "lit_config.json"
     config = {"block_size": 128, "vocab_size": 50, "n_layer": 2, "n_head": 4, "n_embd": 8, "rotary_percentage": 1}
@@ -49,7 +49,7 @@ def test_main(fake_checkpoint_dir, monkeypatch, version, tensor_like):
 
 @pytest.mark.parametrize("version", ("", "_v2"))
 def test_cli(version):
-    cli_path = Path(__file__).parent.parent / "generate" / f"adapter{version}.py"
+    cli_path = Path(__file__).parent.parent / f"litgpt/generate/adapter{version}.py"
     output = subprocess.check_output([sys.executable, cli_path, "-h"])
     output = str(output.decode())
     assert "Generates a response" in output

tests/test_generate_lora.py

@@ -15,7 +15,7 @@
 
 @mock.patch.dict(os.environ, {"LT_ACCELERATOR": "cpu"})
 def test_main(fake_checkpoint_dir, monkeypatch, tensor_like):
-    import generate.lora as generate
+    import litgpt.generate.lora as generate
 
     config_path = fake_checkpoint_dir / "lit_config.json"
     config = {
@@ -56,7 +56,7 @@ def test_main(fake_checkpoint_dir, monkeypatch, tensor_like):
 
 
 def test_cli():
-    cli_path = Path(__file__).parent.parent / "generate" / "lora.py"
+    cli_path = Path(__file__).parent.parent / "litgpt/generate/lora.py"
     output = subprocess.check_output([sys.executable, cli_path, "-h"])
     output = str(output.decode())
     assert "Generates a response" in output