summaryrefslogtreecommitdiffstats
path: root/t3nsor/__init__.py
diff options
context:
space:
mode:
Diffstat (limited to 't3nsor/__init__.py')
-rw-r--r--t3nsor/__init__.py47
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