Update app.py

app.py

@@ -48,7 +48,7 @@ def process_queries(processor, queries, mock_image, max_length: int = 50):
     return batch_query
 
 
-def search(query: str, ds, images) -> str:
+def search(query: str, ds, images):
     qs = []
     with torch.no_grad():
         batch_query = process_queries(processor, [query], mock_image)
@@ -59,9 +59,8 @@ def search(query: str, ds, images) -> str:
     # run evaluation
     retriever_evaluator = CustomEvaluator(is_multi_vector=True)
     scores = retriever_evaluator.evaluate(qs, ds)
-
-    return f"The most relevant page is {scores.argmax(axis=1)}",  images[scores.argmax(axis=1)]
-    # return f"Query: {query}, most relevant page: 1, {len(ds)}", images[1]
+    best_page = int(scores.argmax(axis=1).item())
+    return f"The most relevant page is {best_page}",  images[best_page]
 
 
 def index(file, ds):