A sampling function that returns top token probabilities

[WangHaoranRobin]

I was using playing around with the server example and wanted to expose the probabilities of the generated tokens to the server client to implement custom stopping sequences and criteria(similar to openai's api here).

All it would take should just be creating a different version of "llama_sample_token" and "llama_sample_token_greedy" that returns an object containing the top X tokens and their probabilities.

The only related issue/pr/discussion I was able to find is this pr about logging probabilities. Please give me pointers if similar requests have been discussed somewhere.

Since I'm relatively new to the repo, what is the protocol here? Should I just make a PR?

[nauyisu022]

Did you find out how to achieve this?

[WangHaoranRobin]

Yes. I realized the parameters passed into the sample functions are references. So no need to change the core logic to get the probs. just read directly from the candidate list object after passing it through the sampling function is enough.

i.e. in server.cpp :
https://github.com/ggerganov/llama.cpp/blob/049aa16b8c5c6d086246e4e6b9feb18de4fbd663/examples/server/server.cpp#L293
Read candidates_p.data[i].p after the sampling functions are called to get the probs.

I'm working on a PR for server example about this rn.