thomasht86 commited on
Commit
b62467a
·
verified ·
1 Parent(s): 2a31156

Upload folder using huggingface_hub

Browse files
Files changed (2) hide show
  1. .gitignore +1 -0
  2. main.py +24 -16
.gitignore CHANGED
@@ -8,3 +8,4 @@ template/
8
  *.json
9
  output/
10
  pdfs/
 
 
8
  *.json
9
  output/
10
  pdfs/
11
+ static/saved/
main.py CHANGED
@@ -20,6 +20,7 @@ from backend.colpali import (
20
  is_special_token,
21
  )
22
  from backend.modelmanager import ModelManager
 
23
  from backend.vespa_app import VespaQueryClient
24
  from frontend.app import (
25
  ChatResult,
@@ -82,6 +83,9 @@ But, you should NOT include backticks (`) or HTML tags in your response.
82
  gemini_model = genai.GenerativeModel(
83
  "gemini-1.5-flash-8b", system_instruction=GEMINI_SYSTEM_PROMPT
84
  )
 
 
 
85
 
86
 
87
  @app.on_event("startup")
@@ -102,7 +106,7 @@ def generate_query_id(query):
102
 
103
  @rt("/static/{filepath:path}")
104
  def serve_static(filepath: str):
105
- return FileResponse(f"./static/{filepath}")
106
 
107
 
108
  @rt("/")
@@ -189,6 +193,8 @@ async def get(request, query: str, nn: bool = True):
189
  print(
190
  f"Search results fetched in {end - start:.2f} seconds, Vespa says searchtime was {result['timing']['searchtime']} seconds"
191
  )
 
 
192
  # Start generating the similarity map in the background
193
  asyncio.create_task(
194
  generate_similarity_map(
@@ -269,17 +275,19 @@ async def get_sim_map(query_id: str, idx: int, token: str):
269
 
270
 
271
  async def update_full_image_cache(docid: str, query_id: str, idx: int, image_data: str):
272
- result = result_cache.get(query_id)
273
- if result is None:
274
- await asyncio.sleep(0.5)
275
- return
276
- search_results = get_results_children(result)
277
- # Check if idx exists in list of children
278
- if idx >= len(search_results):
279
- await asyncio.sleep(0.5)
280
- return
281
- search_results[idx]["fields"]["full_image"] = image_data
 
282
  result_cache.set(query_id, result)
 
283
  return
284
 
285
 
@@ -291,12 +299,12 @@ async def full_image(docid: str, query_id: str, idx: int):
291
  image_data = await vespa_app.get_full_image_from_vespa(docid)
292
  # Update the cache with the full image data asynchronously to not block the request
293
  asyncio.create_task(update_full_image_cache(docid, query_id, idx, image_data))
294
- # Decode the base64 image data
295
- # image_data = base64.b64decode(image_data)
296
- image_data = "data:image/jpeg;base64," + image_data
297
-
298
  return Img(
299
- src=image_data,
300
  alt="something",
301
  cls="result-image w-full h-full object-contain",
302
  )
 
20
  is_special_token,
21
  )
22
  from backend.modelmanager import ModelManager
23
+ from pathlib import Path
24
  from backend.vespa_app import VespaQueryClient
25
  from frontend.app import (
26
  ChatResult,
 
83
  gemini_model = genai.GenerativeModel(
84
  "gemini-1.5-flash-8b", system_instruction=GEMINI_SYSTEM_PROMPT
85
  )
86
+ STATIC_DIR = Path(__file__).parent / "static"
87
+ IMG_DIR = STATIC_DIR / "saved"
88
+ os.makedirs(STATIC_DIR, exist_ok=True)
89
 
90
 
91
  @app.on_event("startup")
 
106
 
107
  @rt("/static/{filepath:path}")
108
  def serve_static(filepath: str):
109
+ return FileResponse(STATIC_DIR / filepath)
110
 
111
 
112
  @rt("/")
 
193
  print(
194
  f"Search results fetched in {end - start:.2f} seconds, Vespa says searchtime was {result['timing']['searchtime']} seconds"
195
  )
196
+ # Add result to cache
197
+ result_cache.set(query_id, result)
198
  # Start generating the similarity map in the background
199
  asyncio.create_task(
200
  generate_similarity_map(
 
275
 
276
 
277
  async def update_full_image_cache(docid: str, query_id: str, idx: int, image_data: str):
278
+ result = None
279
+ max_wait = 20 # seconds. If horribly slow network latency.
280
+ start_time = time.time()
281
+ while result is None and time.time() - start_time < max_wait:
282
+ result = result_cache.get(query_id)
283
+ if result is None:
284
+ await asyncio.sleep(0.1)
285
+ try:
286
+ result["root"]["children"][idx]["fields"]["full_image"] = image_data
287
+ except KeyError as err:
288
+ print(f"Error updating full image cache: {err}")
289
  result_cache.set(query_id, result)
290
+ print(f"Full image cache updated for query_id {query_id}")
291
  return
292
 
293
 
 
299
  image_data = await vespa_app.get_full_image_from_vespa(docid)
300
  # Update the cache with the full image data asynchronously to not block the request
301
  asyncio.create_task(update_full_image_cache(docid, query_id, idx, image_data))
302
+ # Save the image to a file
303
+ img_path = IMG_DIR / f"{docid}.jpg"
304
+ with open(img_path, "wb") as f:
305
+ f.write(base64.b64decode(image_data))
306
  return Img(
307
+ src=f"/static/saved/{docid}.jpg",
308
  alt="something",
309
  cls="result-image w-full h-full object-contain",
310
  )