phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-25 18:54:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

(feat) add support for echo for HF logprobs

Browse source
This commit is contained in:
ishaan-jaff 2023-10-31 18:15:17 -07:00
parent 847de09308
commit ce462824be
2 changed files with 16 additions and 2 deletions
Download patch file Download diff file Expand all files Collapse all files

14
litellm/main.py
View file

@ -1792,11 +1792,21 @@ def text_completion(*args, **kwargs):
if kwargs["model"] in litellm.open_ai_text_completion_models and response._hidden_params.get("original_response", None) != None:
return response._hidden_params.get("original_response", None)
transformed_logprobs = None
# only supported for TGI models
try:
raw_response = response._hidden_params.get("original_response", None)
tokens = []
token_logprobs = []
if "prefill" in raw_response[0]["details"]:
prefills = raw_response[0]["details"]['prefill']
for prefill in prefills:
tokens.append(prefill['text'])
token_logprobs.append(prefill['logprob'])
new_tokens = [token['text'] for token in raw_response[0]['details']['tokens']]
new_token_logprobs = [token['logprob'] for token in raw_response[0]['details']['tokens']]
transformed_logprobs = {
"tokens": [token['text'] for token in raw_response[0]['details']['tokens']],
"token_logprobs": [token['logprob'] for token in raw_response[0]['details']['tokens']]
"tokens": tokens + new_tokens,
"token_logprobs": token_logprobs + new_token_logprobs
}
except Exception as e:
print("LiteLLM non blocking exception", e)