diff options
Diffstat (limited to '')
| -rw-r--r-- | t3nsor/__init__.py | 47 |
1 files changed, 23 insertions, 24 deletions
diff --git a/t3nsor/__init__.py b/t3nsor/__init__.py index 5cd12608..9079f9d9 100644 --- a/t3nsor/__init__.py +++ b/t3nsor/__init__.py @@ -21,9 +21,9 @@ class T3nsorResponse: class Usage: def __init__(self, usage_dict: dict) -> None: - self.prompt_tokens = usage_dict['prompt_tokens'] - self.completion_tokens = usage_dict['completion_tokens'] - self.total_tokens = usage_dict['total_tokens'] + self.prompt_tokens = usage_dict['prompt_chars'] + self.completion_tokens = usage_dict['completion_chars'] + self.total_tokens = usage_dict['total_chars'] def __repr__(self): return f'''<__main__.APIResponse.Usage(\n prompt_tokens = {self.prompt_tokens},\n completion_tokens = {self.completion_tokens},\n total_tokens = {self.total_tokens})object at 0x1337>''' @@ -95,24 +95,23 @@ class StreamCompletion: 'prompt' : prompt }) - for resp in response.iter_lines(): - if resp: - yield T3nsorResponse({ - 'id' : f'cmpl-1337-{int(time())}', - 'object' : 'text_completion', - 'created': int(time()), - 'model' : Completion.model, - - 'choices': [{ - 'text' : resp.decode(), - 'index' : 0, - 'logprobs' : None, - 'finish_reason' : 'stop' - }], - - 'usage': { - 'prompt_chars' : len(prompt), - 'completion_chars' : len(resp.decode()), - 'total_chars' : len(prompt) + len(resp.decode()) - } - })
\ No newline at end of file + for chunk in response.iter_content(chunk_size = 2046): + yield T3nsorResponse({ + 'id' : f'cmpl-1337-{int(time())}', + 'object' : 'text_completion', + 'created': int(time()), + 'model' : Completion.model, + + 'choices': [{ + 'text' : chunk.decode(), + 'index' : 0, + 'logprobs' : None, + 'finish_reason' : 'stop' + }], + + 'usage': { + 'prompt_chars' : len(prompt), + 'completion_chars' : len(chunk.decode()), + 'total_chars' : len(prompt) + len(chunk.decode()) + } + })
\ No newline at end of file |