byaldi

app.py

@@ -4,11 +4,20 @@ from datetime import datetime
 
 import dotenv
 import lancedb
+import requests
 from datasets import load_dataset
 from fasthtml.common import *  # noqa
 from huggingface_hub import login, whoami
 from rerankers import Reranker
 
+
+def get_images(query: str):
+    url = "http://147.189.194.113:80/get_pages"
+
+    response = requests.get(url, params={"query": query})
+    return response.json()
+
+
 dotenv.load_dotenv()
 login(token=os.environ.get("HF_TOKEN"))
 
@@ -77,7 +86,7 @@ def _format_results(results):
     return ret
 
 
-def retrieve_and_rerank(query, k=5):
+def retrieve_and_rerank(query, k=3):
     # retrieve ---
     n_fetch = 25
 
@@ -163,6 +172,21 @@ style = Style("""
             margin-right: 10px;
             vertical-align: middle;
         }
+        .image-results {
+            display: flex;
+            flex-wrap: wrap;
+            gap: 10px;
+            margin-top: 20px;
+        }
+        .image-result {
+            width: calc(20% - 10px);
+            text-align: center;
+        }
+        .image-result img {
+            max-width: 100%;
+            height: auto;
+            border-radius: 5px;
+        }
     """)
 
 # get the fast app and route
@@ -229,6 +253,13 @@ def SearchResult(result):
     )
 
 
+def ImageResult(image):
+    return Div(
+        Img(src=f"data:image/jpeg;base64,{image}", alt="arxiv image"),
+        cls="image-result",
+    )
+
+
 def log_query_and_results(query, results):
     log_entry = {
         "timestamp": datetime.now(),
@@ -241,9 +272,19 @@ def log_query_and_results(query, results):
 @rt("/search")
 async def post(query: str):
     results = retrieve_and_rerank(query)
+    image_results = get_images(query)
+
     log_query_and_results(query, results)
 
-    return Div(*[SearchResult(r) for r in results], id="search-results")
+    return Div(
+        H3("Byaldi Results"),
+        Div(*[ImageResult(img) for img in image_results], cls="image-results"),
+        H3("Text Results"),
+        Div(*[SearchResult(r) for r in results], id="text-results"),
+        id="search-results",
+    )
+
+    # return Div(*[SearchResult(r) for r in results], id="search-results")
 
 
 def LogEntry(entry):