Spaces:
Running
on
T4
Running
on
T4
Upload folder using huggingface_hub
Browse files- .gitignore +1 -0
- 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(
|
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 =
|
273 |
-
|
274 |
-
|
275 |
-
|
276 |
-
|
277 |
-
|
278 |
-
|
279 |
-
|
280 |
-
|
281 |
-
|
|
|
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 |
-
#
|
295 |
-
|
296 |
-
|
297 |
-
|
298 |
return Img(
|
299 |
-
src=
|
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 |
)
|