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        stream = None if stream_ptr is None else torch.cuda.ExternalStream(
 30            stream_ptr)
 31        with torch.cuda.stream(stream):
 32            mask = mask.to(logits.device, non_blocking=True)
 33            logits += mask
 34
 35
 36# The recommended way to create a customized batched logits processor:
 37#     * Subclass BatchedLogitsProcessor and implement the processing logics in the __call__ method.
 38#     * Create an instance and pass to LLM.
 39# Alternatively, you can create any callable with the same signature with the __call__ method.
 40# A batched logits processor's arguments for all requests in a batch are made available as lists.
 41# This helps user optimize the callback for large batch sizes. For example:
 42# 1. Process more work on host, e.g. running a JSON state machine, in parallel with model forward pass on device.
 43# 2. Coalesce H2D memory transfers for all requests into a single cudaMemcpyAsync call.
 44# 3. Launch a single batched kernel, e.g. for updating logits on device.
 45class MyBatchedLogitsProcessor(BatchedLogitsProcessor):
 46
 47    def __init__(self, allowed_token_id: int):
 48        self.allowed_token_id = allowed_token_id
 49
 50    def __call__(self, req_ids: List[int], logits: List[torch.Tensor],
 51                 token_ids: List[List[List[int]]], stream_ptr: int,
 52                 client_ids: List[Optional[int]]):
 53        # Generate masks for all requests on host
 54        masks = []
 55        for req_id, req_logits, req_token_ids, client_id in zip(
 56                req_ids, logits, token_ids, client_ids):
 57            mask = torch.full_like(req_logits,
 58                                   fill_value=float("-inf"),
 59                                   device="cpu")
 60            mask[:, :, self.allowed_token_id] = 0
 61            masks.append(mask)
 62
 63        # Move masks to device and add to logits using non-blocking operations
 64        with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)):
 65            for req_logits, mask in zip(logits, masks):
 66                req_logits += mask.to(req_logits.device, non_blocking=True)
 67
 68
 69def main():
 70
 71    # Batched logits processor (only supported in TensorRT backend)
 72    # should be specified when initializing LLM.
 73    llm = LLM(
 74        model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
 75        batched_logits_processor=MyBatchedLogitsProcessor(allowed_token_id=42))
 76
 77    # Sample prompts
 78    prompts = [
 79        "Hello, my name is",
 80        "The president of the United States is",
 81    ]
 82
 83    # Generate text
 84    for prompt_id, prompt in enumerate(prompts):
 85        # Use non-batched logits processor callback only for odd-numbered prompts
 86        if prompt_id % 2 == 0:
 87            sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
 88        else:
 89            # Each prompt can be specified with a logits processor at runtime
 90            sampling_params = SamplingParams(
 91                temperature=0.8,
 92                top_p=0.95,
 93                logits_processor=MyLogitsProcessor(allowed_token_id=42))
 94
 95        for output in llm.generate([prompt], sampling_params):
 96            print(
 97                f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
 98            )
 99
100    # Got output like
101    # 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'
102    # Prompt: 'The president of the United States is', Generated text: "''''''''''''''''''''''''''''''''"
103
104    # Use batched processor with batch size = 2
105    sampling_params = SamplingParams(apply_batched_logits_processor=True)
106    for output in llm.generate(prompts, sampling_params):
107        print(
108            f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
109        )
110
111    # Got output like
112    # Prompt: 'Hello, my name is', Generated text: "''''''''''''''''''''''''''''''''"
113    # Prompt: 'The president of the United States is', Generated text: "''''''''''''''''''''''''''''''''"
114
115
116if __name__ == '__main__':
117    main()