Add vocab_size to Tokenizer.train (Lightning-AI#69)

lit_llama/tokenizer.py

@@ -44,6 +44,6 @@ def decode(self, tokens: torch.Tensor) -> str:
         return self.processor.decode(tokens.tolist())
 
     @staticmethod
-    def train(input: str, destination: str) -> None:
+    def train(input: str, destination: str, vocab_size=32000) -> None:
         model_prefix = os.path.join(destination, "tokenizer")
-        SentencePieceTrainer.Train(input=input, model_prefix=model_prefix)
+        SentencePieceTrainer.Train(input=input, model_prefix=model_prefix, vocab_size=vocab_size)

scripts/prepare_shakespeare.py

@@ -45,7 +45,7 @@ def prepare(destination_path: Path = Path("data/shakespeare")) -> None:
 
     from lit_llama import Tokenizer
 
-    Tokenizer.train(input=input_file_path, destination=destination_path)
+    Tokenizer.train(input=input_file_path, destination=destination_path, vocab_size=100)
     tokenizer = Tokenizer(destination_path / "tokenizer.model")
     train_ids = tokenizer.encode(train_data)
     val_ids = tokenizer.encode(val_data)

train.py

@@ -53,6 +53,7 @@ def main() -> None:
 
     config = LLaMAConfig.from_name("7B")
     config.block_size = block_size
+    config.vocab_size = 100  # from prepare_shakespeare.py
 
     with fabric.device:
         model = LLaMA(config)