Update app.py

app.py

@@ -10,6 +10,8 @@ from PIL import Image
 from torch.utils.data import DataLoader
 from tqdm import tqdm
 
+from functools import partial
+
 from pqdm.processes import pqdm
 
 from colpali_engine.models import ColQwen2, ColQwen2Processor
@@ -42,7 +44,7 @@ DEFAULT_SYSTEM_PROMPT = """
     Answer in the same language as the query.
 """
 
-def query_gpt4o_mini(query, images, api_key, system_prompt=DEFAULT_SYSTEM_PROMPT):
+def query_gpt4o_mini(query, images, api_key=None, system_prompt=DEFAULT_SYSTEM_PROMPT):
     """Calls OpenAI's GPT-4o-mini with the query and image data."""
 
     if api_key and api_key.startswith("sk"):
@@ -135,9 +137,15 @@ def extract_context(images, api_key, window=10):
     prompt = "Give the general context about these pages."
     window_contexts = []
 
-    args = [(prompt, (images[max(i-window+1, 0):i+1], None), api_key, DEFAULT_CONTEXT_PROMPT) 
-            for i in range(0, len(images), window)]
-    window_contexts = pqdm(args, query_gpt4o_mini, n_jobs=8)
+    args = [
+        {
+            'query': prompt,
+            'images': zip(images[max(i-window+1, 0):i+1], [None]*len(images[max(i-window+1, 0):i+1])),
+            'api_key': api_key,
+            'system_prompt': DEFAULT_CONTEXT_PROMPT
+        } for i in range(0, len(images), window)
+    ]
+    window_contexts = pqdm(args, query_gpt4o_mini, n_jobs=8, argument_type='kwargs')
 
     # for i in tqdm(range(0, len(images), window), desc="Extracting context", total=len(images)//window):
     #     window_images = images[max(i-window+1, 0):i+1]
@@ -149,6 +157,8 @@ def extract_context(images, api_key, window=10):
         context = window_contexts[i//window]
         contexts.append(context)
 
+    print(f"Example context: {contexts[0]}")
+
     assert len(contexts) == len(images)
     return contexts