Refactor generate.py for modularity (#906) (#918)

This PR refactors the chat function in generate.py by creating a `Generator` class, removing unnecessary global variables and simplifying the code structure. Tokens and metrics are now yielded by the Generator rather than being printed directly to stdout, making it easier to re-use this code for non-CLI tools.

**Tests:**
Generate
```
python3 torchchat.py generate stories15M --prompt "Once upon a time,"

Using device=mps
Loading model...
Time to load model: 0.46 seconds
-----------------------------------------------------------
Once upon a time, there was a little girl named Lily. One day, she went to the park with her mom. They saw a big tree with lots of pears on it. Lily wanted to eat a pear, but they were too high up. She tried to jump, but she couldn't reach them.
Then, a boy came and took a pear from the tree. Lily was surprised! She thought the boy would be mean, but he was harmless and didn't mean to hurt her. She asked him if she could have a pear too. The boy said yes and gave her a pear. Lily was very happy and thanked the boy for sharing. They sat under the tree and ate the pear together. It was a good day at the park. Once upon a time, there was a little girl named Lily. She loved to play in the garden with her mommy. One day, they were planting some
Time for inference 1: 3.55 sec total, time to first token 0.00 sec with parallel prefill, 199 tokens, 56.01 tokens/sec, 17.85 ms/token
Bandwidth achieved: 2.73 GB/s
*** This first iteration will include cold start effects for dynamic import, hardware caches. ***

========================================

Average tokens/sec: 56.01

========================================

Average tokens/sec: 119.35
Memory used: 0.00 GB
```

Eval
```
python3 torchchat.py eval stories15M --tasks wikitext --limit 10

NumExpr defaulting to 10 threads.
PyTorch version 2.5.0.dev20240629 available.
Using device=mps
Loading model...
Time to load model: 0.38 seconds
-----------------------------------------------------------
Using device 'mps'
[Task: wikitext] metric word_perplexity is defined, but aggregation is not. using default aggregation=weighted_perplexity
[Task: wikitext] metric word_perplexity is defined, but higher_is_better is not. using default higher_is_better=False
[Task: wikitext] metric byte_perplexity is defined, but aggregation is not. using default aggregation=weighted_perplexity
[Task: wikitext] metric byte_perplexity is defined, but higher_is_better is not. using default higher_is_better=False
[Task: wikitext] metric bits_per_byte is defined, but aggregation is not. using default aggregation=bits_per_byte
[Task: wikitext] metric bits_per_byte is defined, but higher_is_better is not. using default higher_is_better=False
Repo card metadata block was not found. Setting CardData to empty.
Repo card metadata block was not found. Setting CardData to empty.
Building contexts for wikitext on rank 0...
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:00<00:00, 541.59it/s]
Running loglikelihood_rolling requests
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 10/10 [00:04<00:00,  2.25it/s]
Time to run eval: 22.57s.
Time in model.forward: 1.01s, over 33 model evaluations
forward run time stats - Median: 0.02s Min: 0.01s Max: 0.27s
For model /Users/puri/.torchchat/model-cache/stories15M/stories15M.pt
wikitext:
 word_perplexity,none: 47350.8811
 byte_perplexity,none: 7.7811
 bits_per_byte,none: 2.9600
 alias: wikitext
```

Co-authored-by: vmpuri <[email protected]>

Author: dbort

eval.py

@@ -4,8 +4,7 @@
 # This source code is licensed under the license found in the
 # LICENSE file in the root directory of this source tree.
 import argparse
-import time
-from typing import Optional
+from typing import Callable, Optional
 
 import torch
 import torch._dynamo.config
@@ -20,7 +19,6 @@
 from build.model import Transformer
 from build.utils import set_precision
 from cli import add_arguments_for_verb, arg_init
-from generate import encode_tokens, model_forward
 from utils.measure_time import measure_time
 
 torch._dynamo.config.automatic_dynamic_shapes = True
@@ -85,11 +83,17 @@ def __init__(
         self,
         model: Transformer,
         tokenizer,
+        model_forward: Optional[Callable] = None,
         max_seq_length: Optional[int] = None,
         device="cpu",
     ):
         super().__init__(device=device)
         self._model = model
+        self._model_forward = (
+            model_forward
+            if model_forward is not None
+            else lambda x, input_pos: model(x, input_pos)
+        )
         self._tokenizer = tokenizer
         self._device = torch.device(device)
         self._max_seq_length = 2048 if max_seq_length is None else max_seq_length
@@ -116,11 +120,8 @@ def device(self):
         return self._device
 
     def tok_encode(self, string: str, **kwargs):
-        encoded = encode_tokens(self._tokenizer, string, bos=True, device=self._device)
-        # encoded is a pytorch tensor, but some internal logic in the
-        # eval harness expects it to be a list instead
-        # TODO: verify this for multi-batch as well
-        encoded = encoded.tolist()
+        bos_id = self._tokenizer.bos_id()
+        encoded = [bos_id] + self._tokenizer.encode(string)
         return encoded
 
     def tok_decode(self, tokens):
@@ -142,7 +143,7 @@ def _model_call(self, inps):
         )
         x = seq.index_select(0, input_pos).view(1, -1)
         with measure_time(message=None) as measure:
-            logits = model_forward(self._model, x, input_pos)
+            logits = self._model_forward(x, input_pos)
         self.times.append(measure.get_time())
         return logits
 
@@ -153,6 +154,7 @@ def _model_generate(self, context, max_length, eos_token_id):
 @torch.no_grad()
 def eval(
     model: Transformer,
+    model_forward: Callable,
     tokenizer,
     tasks: Optional[list] = None,
     limit: Optional[int] = None,
@@ -176,7 +178,11 @@ def eval(
         tasks = ["wikitext"]
 
     model_eval_wrapper = GPTFastEvalWrapper(
-        model, tokenizer, max_seq_length, device=device
+        model,
+        tokenizer,
+        model_forward=model_forward,
+        max_seq_length=max_seq_length,
+        device=device,
     )
 
     try:
@@ -231,11 +237,12 @@ def main(args) -> None:
     )
     tokenizer_args.validate_model(model)
 
+    model_forward = lambda x, input_pos: model(x, input_pos)  # noqa
+
     if compile:
         assert not (
             builder_args.dso_path or builder_args.pte_path
         ), "cannot compile exported model"
-        global model_forward
         model_forward = torch.compile(
             model_forward, mode="reduce-overhead", dynamic=True, fullgraph=True
         )
@@ -244,6 +251,7 @@ def main(args) -> None:
     with measure_time("Time to run eval: {time:.02f}s."):
         result = eval(
             model.to(device),
+            model_forward,
             tokenizer,
             tasks,
             limit,