refactor: add core module

Author: shenxiangzhuang

toyllm/cli/run_gpt2.py

@@ -2,21 +2,22 @@
 
 import typer
 
+from toyllm.core import GenerationConfig
 from toyllm.gpt2 import GPTModel, GPTModelSize, GPTTextGenerator
 
 
 def main(
     prompt: str = "Alan Turing theorized that computers would one day become",
     model_size: GPTModelSize = GPTModelSize.SMALL,
-    max_gen_tokens: int = 40,
+    max_new_tokens: int = 40,
 ) -> None:
     gpt_model = GPTModel(model_size).load()
     text_generator = GPTTextGenerator(gpt_model=gpt_model)
 
     start_time = time.time()
     generate_text = text_generator.generate(
         prompt=prompt,
-        max_gen_tokens=max_gen_tokens,
+        config=GenerationConfig(max_new_tokens=max_new_tokens),
     )
     print(generate_text)
     end_time = time.time()

toyllm/cli/run_speculative_sampling.py

@@ -1,16 +1,19 @@
 import typer
 from rich.console import Console
 
+from toyllm.core import GenerationConfig
 from toyllm.gpt2 import GPTModel, GPTModelSize, GPTTextGenerator, gpt2_tokenizer
 from toyllm.sps import GPTSpsModel, SpsTextGenerator
 from toyllm.util import Timer
 
 
 def main(
     prompt_text: str = "Alan Turing theorized that computers would one day become",
-    generate_tokens: int = 256,
+    max_new_tokens: int = 256,
     k: int = 4,  # K in sps paper
 ) -> None:
+    generate_config = GenerationConfig(max_new_tokens=max_new_tokens)
+
     console = Console()
     console.print(f"Prompt: {prompt_text}", style="bold blue")
     # Test the speculative sampling
@@ -25,8 +28,7 @@ def main(
     with Timer(name="Speculative Sampling"):
         generate_text = sps_text_generator.generate(
             prompt=prompt_text,
-            target_seq_len=generate_tokens,
-            temperature=0,
+            config=generate_config,
         )
     console.print(f"Generated: {generate_text[:200]}", style="bold green")
     console.print(f"{'-' * 20} Speculative Sampling {'-' * 20}", style="bold blue")
@@ -39,7 +41,7 @@ def main(
     with Timer(name="Naive GPT2 Auto-Regressive"):
         generate_text = gpt_text_generator.generate(
             prompt=prompt_text,
-            max_gen_tokens=generate_tokens,
+            config=generate_config,
         )
     console.print(f"Generated: {generate_text[:200]}", style="bold green")
     console.print(f"{'-' * 20} Naive GPT2 Auto-Regressive {'-' * 20}", style="bold blue")

toyllm/gpt2/train.py

@@ -10,6 +10,7 @@
 import torch
 from torch.utils.data import DataLoader
 
+from toyllm.core import GenerationConfig
 from toyllm.device import current_device
 from toyllm.gpt2.config import GPTModelSize, GPTTrainingConfig
 from toyllm.gpt2.dataset import GPTDataloader
@@ -70,7 +71,14 @@ def evaluate_model(
 def generate_and_print_sample(model: GPTModel, tokenizer: tiktoken.Encoding, start_context: str) -> None:
     model.eval()
     text_generate = GPTTextGenerator(gpt_model=model, tokenizer=tokenizer)
-    generate_text = text_generate.generate(prompt=start_context, max_gen_tokens=50, temperature=0.9, top_k=10)
+    generate_text = text_generate.generate(
+        prompt=start_context,
+        config=GenerationConfig(
+            max_new_tokens=50,
+            temperature=0.9,
+            top_k=10,
+        ),
+    )
     print(generate_text.replace("\n", " "))  # Compact print format
     model.train()