added simple generation

Author: jaysonfrancis

test/generate/generation.py

@@ -0,0 +1,104 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional, Tuple
+
+import torch
+
+
+def sample(
+    logits: torch.Tensor,
+    *,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+) -> torch.Tensor:
+    """Sample from a probability distribution
+
+    Args:
+        logits (torch.Tensor): logits from which to sample (vocab_size,)
+        temperature (float): value to scale logits by, default 1.0.
+        top_k (Optional[int]): if specified, prune sampling to only tokens within the top_k probs.
+
+    Returns:
+        torch.Tensor: sampled token id
+    """
+
+    # scale
+    logits = logits / max(temperature, 1e-5)
+
+    # top-k
+    if top_k is not None:
+        v, _ = torch.topk(logits, k=min(top_k, logits.size(-1)))  # (k,)
+        # select last value from top_k above as the pivot
+        pivot = v.select(dim=-1, index=-1).unsqueeze(-1)  # (1,)
+        # mask values smaller than pivot to -inf since these should be pruned
+        logits = torch.where(logits < pivot, -float("Inf"), logits)  # (vocab_size, )
+
+    # normalize
+    probs = torch.nn.functional.softmax(logits, dim=-1)
+
+    return torch.argmax(probs, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def generate_next_token(
+    model,
+    x: torch.Tensor,
+    *,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    logits = model(x)  # (B, T, vocab_size)
+    return (
+        sample(
+            logits[0, -1, :].clone(), temperature=temperature, top_k=top_k
+        ).unsqueeze(-1),
+        logits,
+    )
+
+
+@torch.inference_mode()
+def generate(
+    model,
+    prompt: torch.Tensor,
+    *,
+    max_generated_tokens: int,
+    temperature: float = 1.0,
+    top_k: Optional[int] = None,
+    custom_generate_next_token: Optional[Callable] = None,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """ """
+
+    prompt = prompt.view(1, -1) if prompt.ndim == 1 else prompt
+
+    if custom_generate_next_token is None:
+        _generate_next_token = generate_next_token
+    else:
+        _generate_next_token = custom_generate_next_token
+
+    generated_tokens = prompt.clone()
+
+    tokens, generated_logits = generate_next_token(
+        model,
+        x=prompt,
+        temperature=temperature,
+        top_k=top_k,
+    )
+
+    generated_tokens = torch.cat([generated_tokens, tokens], dim=-1)
+
+    for _ in range(max_generated_tokens - 1):
+        tokens = generated_tokens.clone()
+        tokens, logits = _generate_next_token(
+            model,
+            x=tokens.clone(),
+            temperature=temperature,
+            top_k=top_k,
+        )
+
+        generated_tokens = torch.cat([generated_tokens, tokens], dim=-1)
+        generated_logits = logits
+
+    return generated_tokens, generated_logits

test/generate/test_generate.py

@@ -0,0 +1,138 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import argparse
+from typing import Optional
+
+import torch
+import torch.distributed.checkpoint as dcp
+
+from generation import generate
+from torchtitan import utils
+
+from torchtitan.config_manager import JobConfig
+from torchtitan.datasets import build_tokenizer
+from torchtitan.logging import init_logger, logger
+from torchtitan.models import model_name_to_cls, model_name_to_tokenizer, models_config
+
+
+def example_generate(
+    config_path: str,
+    checkpoint_path: str,
+    prompt: str,
+    *,
+    device: str = "cuda",
+    temperature: float = 1.0,
+    max_generated_tokens: int = 32,
+    top_k: Optional[int] = None,
+):
+    init_logger()
+    color = utils.Color
+
+    # Load configuration from toml file
+    config = JobConfig()
+    config.parse_args([f"--job.config_file={config_path}"])
+    config._validate_config()
+
+    # Load tokenizer and model configuration
+    tokenizer = build_tokenizer(
+        model_name_to_tokenizer[config.model.name], config.model.tokenizer_path
+    )
+    model_cls = model_name_to_cls[config.model.name]
+    model_config = models_config[config.model.name][config.model.flavor]
+    model_config.vocab_size = tokenizer.n_words
+
+    # Load model and checkpoint
+    with torch.device(device):
+        model = model_cls.from_model_args(model_config)
+    state_dict = model.state_dict()
+
+    precompute = False
+    if "freqs_cis" in state_dict:
+        del state_dict["freqs_cis"]
+        precompute = True
+
+    logger.info(f"Loading checkpoint at: {checkpoint_path}")
+    dcp.load(state_dict, checkpoint_id=checkpoint_path)
+
+    # Precompute frequency if required
+    if precompute:
+        model.freqs_cis = model._precompute_freqs_cis().to(args.device)
+
+    # Encode input prompt and generate response
+    input_ids = torch.tensor(
+        tokenizer.encode(prompt, bos=False, eos=False), dtype=torch.long
+    ).to(device)
+    logger.info(f"{color.red}Input tokens: {len(input_ids)}{color.reset}")
+
+    responses, _ = generate(
+        model,
+        input_ids,
+        temperature=temperature,
+        max_generated_tokens=max_generated_tokens,
+        top_k=top_k,
+    )
+    logger.info(
+        f"{color.blue}Output tokens: {len(responses[0]) - len(input_ids)}{color.reset}"
+    )
+
+    response = tokenizer.decode(
+        [token.item() for token in responses[0][len(input_ids) :]]
+    )
+    logger.info(f"{color.red}{prompt}{color.blue}{response}")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Test generation")
+    parser.add_argument(
+        "--config", type=str, required=True, help="TOML config file path (required)"
+    )
+    parser.add_argument(
+        "--checkpoint",
+        type=str,
+        required=True,
+        help="Checkpoint path to load (required)",
+    )
+    parser.add_argument(
+        "--device",
+        type=str,
+        default="cuda",
+        choices=["cpu", "cuda"],
+        help="Device to load model on. Default is 'cuda'",
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=1.0,
+        help="Sampling temperature. Default is 1.0",
+    )
+    parser.add_argument(
+        "--max_generated_tokens",
+        type=int,
+        default=32,
+        help="Max number of tokens to generate. Default is 32",
+    )
+    parser.add_argument(
+        "--top_k", type=int, help="Prune to select from top_k probabilities. Optional"
+    )
+    parser.add_argument(
+        "--prompt",
+        type=str,
+        default="Hello! How are you?",
+        help="Input prompt for generation",
+    )
+
+    args = parser.parse_args()
+
+    example_generate(
+        config_path=args.config,
+        checkpoint_path=args.checkpoint,
+        prompt=args.prompt,
+        device=args.device,
+        temperature=args.temperature,
+        max_generated_tokens=args.max_generated_tokens,
+        top_k=args.top_k,
+    )