Update app.py

app.py

@@ -78,14 +78,11 @@ 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)
-
-    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
-
-    results = [{"page": idx, "image": "image_placeholder"} for idx in top_k_indices]
-
-    return {"results": results}
+    best_page = int(scores.argmax(axis=1).item())
+    return f"The most relevant page is {best_page}", images[best_page]
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)