From 19c553549b4a805f167e07c136ff00d64ed11fe2 Mon Sep 17 00:00:00 2001 From: Emmett McFaralne Date: Fri, 17 May 2024 01:02:17 -0400 Subject: [PATCH] fixed token count mismatch --- setup.py | 2 +- thepipe_api/compressor.py | 10 +++++----- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/setup.py b/setup.py index 31a7cb3..78f02cc 100644 --- a/setup.py +++ b/setup.py @@ -2,7 +2,7 @@ setup( name='thepipe_api', - version='0.3.5', + version='0.3.6', author='Emmett McFarlane', author_email='emmett@thepi.pe', description='Automate information extraction for multimodal LLMs.', diff --git a/thepipe_api/compressor.py b/thepipe_api/compressor.py index 2b0738b..f92563f 100644 --- a/thepipe_api/compressor.py +++ b/thepipe_api/compressor.py @@ -4,8 +4,7 @@ import tempfile from typing import List, Optional import os -from .core import Chunk, SourceTypes, print_status, count_tokens -from .thepipe import count_tokens +from .core import Chunk, SourceTypes, print_status from PIL import Image CTAGS_EXECUTABLE_PATH = "C:\ctags.exe" if os.name == 'nt' else "ctags-universal" @@ -97,9 +96,10 @@ def calculate_tokens(chunk: Chunk) -> int: def compress_chunks(chunks: List[Chunk], verbose: bool = False, limit: Optional[int] = None) -> List[Chunk]: new_chunks = chunks for _ in range(min(MAX_COMPRESSION_ATTEMPTS, len(chunks))): - if count_tokens(new_chunks) <= limit: + token_count = sum(calculate_tokens(chk) for chk in new_chunks) + if token_count <= limit: break - if verbose: print_status(f"Compressing prompt ({count_tokens(chunks)} tokens / {limit} limit)", status='info') + if verbose: print_status(f"Compressing prompt ({token_count} tokens / {limit} limit)", status='info') new_chunks = [] chunk_with_most_tokens = max(chunks, key=calculate_tokens) for chunk in chunks: @@ -123,6 +123,6 @@ def compress_chunks(chunks: List[Chunk], verbose: bool = False, limit: Optional[ new_res = (new_chunk.image.width//2, new_chunk.image.height//2) new_chunk.image = new_chunk.image.resize(new_res) new_chunks.append(new_chunk) - if count_tokens(new_chunks) > limit and verbose: + if sum(calculate_tokens(chk) for chk in new_chunks) > limit and verbose: print_status("Failed to compress within limit, continuing", status='error') return new_chunks