Merge branch 'TablewareBox:main' into wangcx

Author: boliqq07

evals/completion_fns/gemini.py

@@ -1,4 +1,6 @@
 import os
+
+import PIL
 import numpy as np
 from typing import Any, Optional, Union
 
@@ -56,7 +58,9 @@ def truncate_multimodal_prompt(prompt_list, max_images=16, max_size_bytes=4 * 10
             # The image data is a string representation of bytes; calculate its length accordingly.
             item_size = len(item['data'])  # Approximation of size in bytes
             image_count += 1
-        else:
+        elif isinstance(item, PIL.Image.Image):  # It's an image
+            item_size = len(item.tobytes())  # Approximation of size in bytes
+            image_count += 1
             continue  # Skip any item that doesn't fit expected structure
 
         # Check if adding this item would exceed limits
@@ -159,6 +163,7 @@ def __call__(
                                       generation_config=generation_config,
                                       safety_settings=safety_settings)
         # response = request_with_timeout(model.generate_content, contents=[openai_create_prompt] + attached_file_content)
+        print("Num tokens:", model.count_tokens(contents=contents))
         response = model.generate_content(
             contents=contents,
         )

evals/utils/doc_utils.py

@@ -1,5 +1,9 @@
+import io
 import os
 from pathlib import Path
+from typing import Union
+
+import PIL
 import fitz  # PyMuPDF
 import PyPDF2
 import tiktoken
@@ -109,7 +113,8 @@ def extract_text_and_fill_in_images(pdf_path, save_img_path=None, add_page_num:
         all_contents.append(text)
         for i, block in enumerate(blocks_from_pymupdf[page_num]):
             if block[1] == "image":
-                img_cookie = {'mime_type': 'image/png', 'data': block[2]}
+                # img_cookie = {'mime_type': 'image/png', 'data': block[2]}
+                img_cookie = PIL.Image.open(io.BytesIO(block[2]))
                 all_contents.append(img_cookie)
                 # if i != len(blocks_from_pymupdf) - 1:
                 #     fig_caption = " ".join(block[i+1].split()[:3])
@@ -124,3 +129,15 @@ def num_tokens_from_string(string: str, encoding_name: str) -> int:
     encoding = tiktoken.get_encoding(encoding_name)
     num_tokens = len(encoding.encode(string))
     return num_tokens
+
+
+def prompt_cutoff(prompt: Union[str, list], model: str = "gpt-3.5-turbo"):
+    model_max_tokens = {
+        "gpt-3.5-turbo": 15800,
+        "claude-2.1": 10000,
+        "gemini-1.0-pro": 25000,
+        "gemini-1.0-pro-vision": 25000,
+    }
+
+    if type(prompt) == str:
+        pass