Update app.py

app.py

@@ -74,28 +74,7 @@ async def index(files: List[UploadFile] = File(...)):
     
     return {"message": f"Uploaded and converted {len(images)} pages"}
 
-@app.get("/search")
-async 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 = []
-    for idx in top_k_indices:
-        img_byte_arr = BytesIO()
-        images[idx].save(img_byte_arr, format='PNG')
-        img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
-        results.append({"image": img_base64, "page": f"Page {idx}"})
-
-    # Generate PDF
+def generate_pdf(results):
     pdf_buffer = BytesIO()
     c = canvas.Canvas(pdf_buffer, pagesize=letter)
     width, height = letter
@@ -118,10 +97,78 @@ async def search(query: str, k: int):
 
     c.save()
     pdf_buffer.seek(0)
+    return pdf_buffer
+
+@app.get("/search")
+async def search(query: str, k: int = 1):
+    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 = []
+    for idx in top_k_indices:
+        img_byte_arr = BytesIO()
+        images[idx].save(img_byte_arr, format='PNG')
+        img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        results.append({"image": img_base64, "page": f"Page {idx}"})
+
+    pdf_buffer = generate_pdf(results)
 
     # Use StreamingResponse to handle in-memory file
     response = StreamingResponse(pdf_buffer, media_type='application/pdf')
-    response.headers['Content-Disposition'] = 'attachment; filename="search_results.pdf"'
+    response.headers['Content-Disposition'] = 'attachment; filename="results.pdf"'
+
+    return response
+
+@app.get("/search_by_cv")
+async def search_by_cv(file: UploadFile = File(...), k: int = 10):
+    # Lire le fichier PDF uploadé
+    content = await file.read()
+    pdf_image_list = convert_from_bytes(content)
+
+    # Générer les embeddings pour les pages du PDF uploadé
+    qs = []
+    dataloader = DataLoader(
+        pdf_image_list,
+        batch_size=4,
+        shuffle=False,
+        collate_fn=lambda x: process_images(processor, x),
+    )
+    for batch_query in dataloader:
+        with torch.no_grad():
+            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"))))
+
+    # Comparer les embeddings du CV uploadé avec ceux déjà indexés
+    retriever_evaluator = CustomEvaluator(is_multi_vector=True)
+    scores = retriever_evaluator.evaluate(qs, ds)
+
+    # Trouver les indices des résultats les plus pertinents
+    top_k_indices = scores.argsort(axis=1)[0][-k:][::-1]
+
+    # Préparer les résultats sous forme d'images
+    results = []
+    for idx in top_k_indices:
+        img_byte_arr = BytesIO()
+        images[idx].save(img_byte_arr, format='PNG')
+        img_base64 = base64.b64encode(img_byte_arr.getvalue()).decode('utf-8')
+        results.append({"image": img_base64, "page": f"Page {idx}"})
+
+    # Générer le PDF des résultats
+    pdf_buffer = generate_pdf(results)
+
+    # Utiliser StreamingResponse pour renvoyer le fichier PDF généré
+    response = StreamingResponse(pdf_buffer, media_type='application/pdf')
+    response.headers['Content-Disposition'] = 'attachment; filename="results.pdf"'
 
     return response