Update app.py

app.py

@@ -42,6 +42,48 @@ mock_image = Image.new("RGB", (448, 448), (255, 255, 255))
 ds = []
 images = []
 
+@app.post("/index")
+def index(files: List[UploadFile] = File(...)):
+    global ds, images
+    images = []
+    ds = []
+    for file in files:
+        content = file.read()
+        pdf_image_list = convert_from_path(io.BytesIO(content))
+        images.extend(pdf_image_list)
+    
+    dataloader = DataLoader(
+        images,
+        batch_size=4,
+        shuffle=False,
+        collate_fn=lambda x: process_images(processor, x),
+    )
+    for batch_doc in dataloader:
+        with torch.no_grad():
+            batch_doc = {k: v.to(device) for k, v in batch_doc.items()}
+            embeddings_doc = model(**batch_doc)
+        ds.extend(list(torch.unbind(embeddings_doc.to("cpu"))))
+    
+    return {"message": f"Uploaded and converted {len(images)} pages"}
+
+@app.post("/search")
+def search(query: str, k: int):
+    qs = []
+    with torch.no_grad():
+        batch_query = process_queries(processor, [query], mock_image)
+        batch_query = {k: v.to(device) for k, v in batch_query.items()}
+        embeddings_query = model(**batch_query)
+        qs.extend(list(torch.unbind(embeddings_query.to("cpu"))))
+
+    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}
+
 # Rediriger la racine vers /docs
 @app.get("/")
 def read_root():