Update app.py

app.py

@@ -78,11 +78,16 @@ async def search(query: str, k: int):
         embeddings_query = model(**batch_query)
         qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
 
-    # run evaluation
     retriever_evaluator = CustomEvaluator(is_multi_vector=True)
     scores = retriever_evaluator.evaluate(qs, ds)
-    best_page = int(scores.argmax(axis=1).item())
-    return f"The most relevant page is {best_page}", images[best_page]
+
+    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
+
+    results = []
+    for idx in top_k_indices:
+        results.append({"image": str(images[idx]), "page": f"Page {idx}"})
+
+    return {"results": results}
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)