Update app.py

app.py

@@ -41,47 +41,13 @@ 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"}
+# Initialiser le pipeline de génération de texte
+pipe = pipeline("text2text-generation", model="google/flan-t5-small")
 
-@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}
+@app.get("/generate")
+def generate(text: str):
+    output = pipe(text)
+    return {"output": output[0]["generated_text"]}
 
 # Rediriger la racine vers /docs
 @app.get("/")