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()