phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

fix(utils.py): ensure streaming output parsing only applied for hf / sagemaker models

Browse source 
selectively applies the <s>
</s> checking
This commit is contained in:
Krrish Dholakia 2024-04-17 17:43:41 -07:00
parent 53df916f69
commit 7d0086d742
2 changed files with 23 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

9
litellm/utils.py
View file

@ -8856,7 +8856,16 @@ class CustomStreamWrapper:
raise e
def check_special_tokens(self, chunk: str, finish_reason: Optional[str]):
"""
Output parse <s> / </s> special tokens for sagemaker + hf streaming.
"""
hold = False
if (
self.custom_llm_provider != "huggingface"
and self.custom_llm_provider != "sagemaker"
):
return hold, chunk
if finish_reason:
for token in self.special_tokens:
if token in chunk: