Control generated text using logits processor#

Source NVIDIA/TensorRT-LLM.

  1### Control generated text using logits processor
  2from typing import List, Optional
  3
  4import torch
  5
  6from tensorrt_llm import LLM
  7from tensorrt_llm.sampling_params import (BatchedLogitsProcessor,
  8                                          LogitsProcessor, SamplingParams)
  9
 10
 11# The recommended way to create a customized logits processor:
 12#     * Subclass LogitsProcessor and implement the processing logics in the __call__ method.
 13#     * Create an instance and pass to SamplingParams.
 14# Alternatively, you can create any callable with the same signature with the __call__ method.
 15# This simple callback will output a specific token at each step irrespective of prompt.
 16# Refer to ../bindings/executor/example_logits_processor.py for a more
 17# sophisticated callback that generates JSON structured output.
 18class MyLogitsProcessor(LogitsProcessor):
 19
 20    def __init__(self, allowed_token_id: int):
 21        self.allowed_token_id = allowed_token_id
 22
 23    def __call__(self, req_id: int, logits: torch.Tensor,
 24                 token_ids: List[List[int]], stream_ptr: int,
 25                 client_id: Optional[int]):
 26        mask = torch.full_like(logits, fill_value=float("-inf"), device="cpu")
 27        mask[:, :, self.allowed_token_id] = 0
 28
 29        with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)):
 30            mask = mask.to(logits.device, non_blocking=True)
 31            logits += mask
 32
 33
 34# The recommended way to create a customized batched logits processor:
 35#     * Subclass BatchedLogitsProcessor and implement the processing logics in the __call__ method.
 36#     * Create an instance and pass to LLM.
 37# Alternatively, you can create any callable with the same signature with the __call__ method.
 38# A batched logits processor's arguments for all requests in a batch are made available as lists.
 39# This helps user optimize the callback for large batch sizes. For example:
 40# 1. Process more work on host, e.g. running a JSON state machine, in parallel with model forward pass on device.
 41# 2. Coalesce H2D memory transfers for all requests into a single cudaMemcpyAsync call.
 42# 3. Launch a single batched kernel, e.g. for updating logits on device.
 43class MyBatchedLogitsProcessor(BatchedLogitsProcessor):
 44
 45    def __init__(self, allowed_token_id: int):
 46        self.allowed_token_id = allowed_token_id
 47
 48    def __call__(self, req_ids: List[int], logits: List[torch.Tensor],
 49                 token_ids: List[List[List[int]]], stream_ptr: int,
 50                 client_ids: List[Optional[int]]):
 51        # Generate masks for all requests on host
 52        masks = []
 53        for req_id, req_logits, req_token_ids, client_id in zip(
 54                req_ids, logits, token_ids, client_ids):
 55            mask = torch.full_like(req_logits,
 56                                   fill_value=float("-inf"),
 57                                   device="cpu")
 58            mask[:, :, self.allowed_token_id] = 0
 59            masks.append(mask)
 60
 61        # Move masks to device and add to logits using non-blocking operations
 62        with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)):
 63            for req_logits, mask in zip(logits, masks):
 64                req_logits += mask.to(req_logits.device, non_blocking=True)
 65
 66
 67def main():
 68
 69    # Batched logits processor should be specified when initializing LLM.
 70    llm = LLM(
 71        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
 72        batched_logits_processor=MyBatchedLogitsProcessor(allowed_token_id=42))
 73
 74    # Sample prompts
 75    prompts = [
 76        "Hello, my name is",
 77        "The president of the United States is",
 78    ]
 79
 80    # Generate text
 81    for prompt_id, prompt in enumerate(prompts):
 82        # Use non-batched logits processor callback only for odd-numbered prompts
 83        if prompt_id % 2 == 0:
 84            sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 85        else:
 86            # Each prompt can be specified with a logits processor at runtime
 87            sampling_params = SamplingParams(
 88                temperature=0.8,
 89                top_p=0.95,
 90                logits_processor=MyLogitsProcessor(allowed_token_id=42))
 91
 92        for output in llm.generate([prompt], sampling_params):
 93            print(
 94                f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
 95            )
 96
 97    # Got output like
 98    # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming'
 99    # Prompt: 'The president of the United States is', Generated text: "''''''''''''''''''''''''''''''''"
100
101    # Use batched processor with batch size = 2
102    sampling_params = SamplingParams(apply_batched_logits_processor=True)
103    for output in llm.generate(prompts, sampling_params):
104        print(
105            f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
106        )
107
108    # Got output like
109    # Prompt: 'Hello, my name is', Generated text: "''''''''''''''''''''''''''''''''"
110    # Prompt: 'The president of the United States is', Generated text: "''''''''''''''''''''''''''''''''"
111
112
113if __name__ == '__main__':
114    main()