LLM Logits Processor

Source https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/llm-api/llm_logits_processor.py.

 1### Control generated text using logits post processor
 2import typing as tp
 3
 4import torch
 5
 6from tensorrt_llm import LLM, SamplingParams
 7
 8
 9# Define the logits post-processor callback. This simple callback will output
10# a specific token at each step irrespective of prompt.
11# Refer to ../bindings/executor/example_logits_processor.py for a more
12# sophisticated callback that generates JSON structured output.
13def logits_post_processor(req_id: int, logits: torch.Tensor,
14                          ids: tp.List[tp.List[int]], stream_ptr: int,
15                          client_id: tp.Optional[int]):
16    target_token_id = 42
17    with torch.cuda.stream(torch.cuda.ExternalStream(stream_ptr)):
18        logits[:] = float("-inf")
19        logits[..., target_token_id] = 0
20
21
22# Several callbacks can be specified when initializing LLM
23llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0",
24          logits_post_processor_map={"my_logits_pp": logits_post_processor})
25
26# Sample prompts
27prompts = [
28    "Hello, my name is",
29    "The president of the United States is",
30]
31
32# Generate text
33for prompt_id, prompt in enumerate(prompts):
34    # We will use logits post processor callback only for odd-numbered prompts
35    if prompt_id % 2 == 0:
36        sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
37    else:
38        # Each prompt can use one callback from the choices that were provided to LLM
39        sampling_params = SamplingParams(
40            temperature=0.8,
41            top_p=0.95,
42            logits_post_processor_name='my_logits_pp')
43
44    for output in llm.generate([prompt], sampling_params):
45        print(
46            f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}"
47        )
48
49# Got output like
50# 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'
51# Prompt: 'The president of the United States is', Generated text: "''''''''''''''''''''''''''''''''"