Question about `input_token_logprobs`

[WissamAntoun]

Hey, thanks a lot for the great work and for releasing the codebase.

I've been trying to create my own annotation on my own dataset using prompt_classify.py but i've been getting errors here:

https://github.com/CodeCreator/WebOrganizer/blob/main/define_domains/prompt_classify.py#L132

assert all(
                len(answer_tokens) > 1
                for answer_tokens in meta_info["input_token_logprobs"]
            ), f"All answers should have at least 2 tokens in {meta_info['input_token_logprobs']}"

The thing is the SGLang classifier has been returning the answer/choice token as the first token. Like the following:

input_token_logprobs = [
[[-4.267991065979004, 32, 'A']],
[[-4.017991065979004, 33, 'B']],
[[-4.517991065979004, 34, 'C']],
[[-4.392991065979004, 35, 'D']],
[[-3.142991065979004, 36, 'E']],
[[-4.517991065979004, 37, 'F']],
[[-3.017991065979004, 38, 'G']],
[[-3.142991065979004, 39, 'H']],
[[-2.392991065979004, 40, 'I']],
[[-3.5137503147125244, 41, 'J']],
[[-3.2658731937408447, 42, 'K']],
[[-3.2658731937408447, 43, 'L']],
[[-2.3887503147125244, 44, 'M']],
[[-3.8887503147125244, 45, 'N']],
[[-4.013750076293945, 46, 'O']],
[[-3.6387503147125244, 47, 'P']],
[[-2.7637503147125244, 48, 'Q']],
[[-2.6387503147125244, 49, 'R']],
[[-2.5137503147125244, 50, 'S']],
[[-2.3887503147125244, 51, 'T']],
[[-4.138750076293945, 52, 'U']],
[[-3.8887503147125244, 53, 'V']],
[[-2.8887503147125244, 54, 'W']],
[[-3.2637503147125244, 55, 'X']]
]

I did some digging in SgLang's code and i found that in their runtime, they actualy remove the first non choice token and recompute the logprobs. Check the code here: https://github.com/sgl-project/sglang/blob/main/python/sglang/lang/backend/runtime_endpoint.py#L267 .

        # Remove extra token if no token healing occurred
        for i in range(len(input_token_logprobs)):
            healed_token_str = input_token_logprobs[i][0][-1]
            if s.text_.endswith(healed_token_str):
                healed_token_logprob = input_token_logprobs[i][0][0]
                normalized_prompt_logprobs[i] = (
                    normalized_prompt_logprobs[i] * len(input_token_logprobs[i])
                    - healed_token_logprob
                ) / (len(input_token_logprobs[i]) - 1)
                input_token_logprobs[i] = input_token_logprobs[i][1:]

Could it be that this they have already address the issue you are accounting for?

[CodeCreator]

Hey! Great to hear you're using the code! SGLang has been under a lot of development, so when I wrote this I was using a version sglang==0.3.0. I suspect it's due to the version mismatch.

If you can open a pull-request that fixes this in your sglang version, that would be awesome and I'd be happy to review and merge! Otherwise, maybe try using sglang==0.3.0 or I might upgrade the code in the next few days.

[WissamAntoun]

Thanks for the reply, I'll test sglang 0.3.0 and I'll report back if it was the cause of the issue. Thanks

[vishaal27]

Hi @WissamAntoun, did using sglang 0.3.0 work for you?

[WissamAntoun]

Well it should probably work, but didn't have time to test it. I ended up just skipping the recomputation and get the choice loss directly.

           demonstration_permutation = get_demonstration_permutation(
                indices[i], prompt_config
            )
            permutation = get_permutation(indices[i], prompt_config)

            # We re-compute answer logprobs, as the first token is the preceding token
            # that is the same for all answers
            # permuted_choice_loss = np.array(
            #     [
            #         sum(logprob for logprob, token_id, _ in answer_tokens[1:])
            #         / (len(answer_tokens) - 1)
            #         for answer_tokens in meta_info["input_token_logprobs"]
            #     ]
            # )
            permuted_choice_loss = np.array(
                [
                    answer_tokens[0][0]
                    for answer_tokens in meta_info["input_token_logprobs"]
                ]
            )
            choice_loss = np.zeros_like(permuted_choice_loss)
            choice_loss[permutation] = permuted_choice_loss