Get Logprobs

In PyTRIO, there are two common ways to get token-level logprobs:

Method	Best for	What you get
`sample()`	Continue generation and record one sampling run	prompt-token logprobs and generated-token logprobs
`compute_logprobs()`	Rescore an already fixed text sequence with a selected model	logprobs for every actual token in the full input text

In short, sample() answers: "During this sampling run, what were the logprobs of the prompt tokens and generated tokens?" compute_logprobs() answers: "Given this complete text, how much does the selected model endorse each token in it?"

Method 1: Get Logprobs With sample

sample() asks the model to continue from a prompt. By default, generated-token logprobs are returned in sequences[*].logprobs. If you set include_prompt_logprobs=True, the response also includes prompt-token logprobs.

import pytrio as trio


client = trio.ServiceClient()
sampling_client = client.create_sampling_client(base_model="Qwen/Qwen3.5-4B")
tokenizer = sampling_client.get_tokenizer()

messages = [
    {"role": "user", "content": "1 + 1 等于多少？"},
]
prompt_text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
    add_generation_prompt=True,
)
prompt_ids = tokenizer.encode(prompt_text, add_special_tokens=False)

response = sampling_client.sample(
    prompt=trio.ModelInput.from_ints(prompt_ids),
    sampling_params=trio.SamplingParams(max_tokens=8, temperature=0.7),
    include_prompt_logprobs=True,
).result()

sequence = response.sequences[0]

prompt_logprobs = response.prompt_logprobs
completion_tokens = list(sequence.tokens)
completion_logprobs = list(sequence.logprobs)

print(type(prompt_logprobs), len(prompt_logprobs))
print(type(completion_logprobs), len(completion_logprobs))
print(completion_tokens)
print(completion_logprobs)
print(sequence.text)

The two directly usable values are lists:

response.prompt_logprobs: list[float | None], aligned with prompt_ids.
sequence.logprobs: list[float | None], aligned with sequence.tokens, the tokens generated in this sampling run.

You can think of the returned structure as:

response.prompt_logprobs
[
    None,
    -0.20,
    -0.00,
    ...
]

response.sequences[0].tokens
[17, 151645]

response.sequences[0].logprobs
[-0.42, -0.02]

response.sequences[0].text
'2'

For a more readable object, zip the generated token IDs and logprobs together:

completion = [
    {
        "token_id": token_id,
        "token": tokenizer.decode([token_id]),
        "logprob": logprob,
    }
    for token_id, logprob in zip(sequence.tokens, sequence.logprobs)
]

The result looks like:

[
    {"token_id": 17, "token": "2", "logprob": -0.42},
    {"token_id": 151645, "token": "<|im_end|>", "logprob": -0.02},
]

The token IDs and values above only show the output shape. Real logprobs depend on the model, chat template, sampling parameters, and generated content. The important part is that both prompt logprobs and generated-token logprobs are directly usable lists for RL, rescoring, or logging.

Method 2: Get Full-Text Logprobs With compute_logprobs

compute_logprobs() does not generate new tokens. It runs a forward scoring pass over the complete ModelInput you provide:

def compute_logprobs(self, prompt: ModelInput) -> APIFuture[list[float | None]]

The returned list is aligned with the input tokens:

The value at position i is the log probability of token i conditioned on the previous tokens.
The first token has no previous context, so it may be None.
The return value is the log probability of each actual token, not a full-vocabulary probability distribution.

import pytrio as trio


client = trio.ServiceClient()
sampling_client = client.create_sampling_client(base_model="Qwen/Qwen3.5-4B")
tokenizer = sampling_client.get_tokenizer()

messages = [
    {"role": "user", "content": "1 + 1 等于多少？"},
    {"role": "assistant", "content": "2"},
]
text = tokenizer.apply_chat_template(
    messages,
    tokenize=False,
)
tokens = tokenizer.encode(text, add_special_tokens=False)

logprobs = sampling_client.compute_logprobs(
    prompt=trio.ModelInput.from_ints(tokens),
).result()

scored_tokens = [
    {
        "token_id": token_id,
        "token": tokenizer.decode([token_id]),
        "logprob": logprob,
    }
    for token_id, logprob in zip(tokens, logprobs)
]

compute_logprobs() returns list[float | None], aligned with the input tokens. The scored_tokens object above looks like:

[
    {"token_id": 151644, "token": "<|im_start|>", "logprob": None},
    {"token_id": 872, "token": "user", "logprob": -0.20},
    {"token_id": 198, "token": "\n", "logprob": -0.00},
    {"token_id": 16, "token": "1", "logprob": -2.18},
    {"token_id": 488, "token": " +", "logprob": -5.86},
    {"token_id": 220, "token": " 1", "logprob": -1.09},
    {"token_id": 151645, "token": "<|im_end|>", "logprob": -0.00},
    {"token_id": 198, "token": "\n", "logprob": -0.00},
    {"token_id": 151644, "token": "<|im_start|>", "logprob": 0.0},
    {"token_id": 77091, "token": "assistant", "logprob": -0.01},
    {"token_id": 198, "token": "\n", "logprob": -0.00},
    {"token_id": 17, "token": "2", "logprob": -7.80},
    {"token_id": 151645, "token": "<|im_end|>", "logprob": -0.02},
]

This is the core difference between compute_logprobs() and sample(): sample() returns prompt and generated-token logprobs for one generation run, while compute_logprobs() can rescore an already fixed full text sequence.

Teacher Logprobs In OPD

In OPD, short for on-policy distillation, the student first generates a completion with the current policy. The teacher then computes logprobs for the exact token trajectory generated by the student. This gives token-level student_logprobs - teacher_logprobs, which can be converted into advantages for importance_sampling.

In OPD and similar model-distillation workflows, the teacher is usually a stronger, more stable model, or a checkpoint that already performs better on the target task. Teacher logprobs are not used to continue generation. They answer a training question: how much does the stronger teacher endorse each token that the student actually sampled? That token-level signal then enters the reverse-KL or advantage calculation and guides the student toward the teacher-preferred token distribution.

The key detail is that the teacher must score prompt + completion, and then you slice out only the completion span.

import numpy as np

import pytrio as trio


def completion_teacher_logprobs(
    teacher_client,
    prompt_ids: list[int],
    completion_ids: list[int],
) -> list[float]:
    all_ids = prompt_ids + completion_ids

    all_logprobs = teacher_client.compute_logprobs(
        prompt=trio.ModelInput.from_ints(all_ids),
    ).result()

    completion_logprobs = all_logprobs[len(prompt_ids):]
    if len(completion_logprobs) != len(completion_ids):
        raise ValueError("teacher logprobs and completion tokens are not aligned")
    if any(value is None for value in completion_logprobs):
        raise ValueError("completion logprobs should not contain None")

    return [float(value) for value in completion_logprobs]


# seq.logprobs are the old-policy logprobs for the completion tokens returned by student sampling.
student_logprobs = [float(value) for value in seq.logprobs]
teacher_logprobs = completion_teacher_logprobs(
    teacher_client=teacher_client,
    prompt_ids=prompt_ids,
    completion_ids=list(seq.tokens),
)

reverse_kl = np.asarray(student_logprobs) - np.asarray(teacher_logprobs)
advantages = -kl_penalty_coef * reverse_kl

The core OPD signal is:

reverse_kl = student_logprobs - teacher_logprobs
advantages = -kl_penalty_coef * reverse_kl

After that, shift and align completion_ids, the student sampling logprobs, and the computed advantages, put them into Datum.loss_fn_inputs, and train with loss_fn="importance_sampling".

Difference From Common Generation APIs

Many cloud generation APIs return only final text by default. Some APIs can also return token logprobs from the generation step. For example, DeepSeek's logprobs parameter returns log probabilities for output tokens, and top_logprobs returns the top-N tokens and log probabilities at each output position. OpenAI Chat Completions and Gemini expose similar logprobs for generated output tokens.

Those APIs usually answer: "What were the probabilities of the tokens generated in this request?" compute_logprobs answers a different question: given an already fixed token sequence, what are the logprobs of the actual tokens in that sequence under the selected model? That is why OPD can let the student sample first and then ask the teacher to rescore the same full trajectory.

Token-level subtraction requires token alignment. In practice, use the same tokenizer for student and teacher, or at least compatible tokenizers. If tokenization differs, student_logprobs - teacher_logprobs cannot be subtracted position by position.

Method 1: Get Logprobs With sample

Method 2: Get Full-Text Logprobs With compute_logprobs

Teacher Logprobs In OPD

Difference From Common Generation APIs

On this page