testRAG / app.py
goldrode's picture
Update app.py
e3af4f4 verified
from fastapi import FastAPI, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from sentence_transformers import SentenceTransformer
import fitz # PyMuPDF pour extraction du texte PDF
import faiss
import os
import numpy as np
import requests
import gradio as gr
# Configuration de l'API Gemini
GEMINI_API_KEY = "AIzaSyArbgg_p_HlmpgrcjVYemdSJeMCP9OTj3E"
GEMINI_API_URL = "https://generativelanguage.googleapis.com/v1beta/models/gemini-1.5-flash-latest:generateContent"
# Configuration des embeddings et FAISS
EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
model = SentenceTransformer(EMBEDDING_MODEL)
INDEX_PATH = "medical_faiss_index"
documents = [] # Liste pour stocker les textes indexés
# Chargement ou création de l'index FAISS
if os.path.exists(INDEX_PATH):
index = faiss.read_index(INDEX_PATH)
print("Index FAISS chargé avec succès.")
else:
index = faiss.IndexFlatL2(model.get_sentence_embedding_dimension())
print("Nouvel index FAISS créé.")
# Fonction pour extraire le texte des fichiers PDF
def extract_text_from_pdf(file_content: bytes):
pdf = fitz.open(stream=file_content, filetype="pdf")
paragraphs = []
for page in pdf:
text = page.get_text()
if text.strip():
paragraphs.extend([p.strip() for p in text.split("\n\n") if p.strip()])
return paragraphs
# Ajouter des documents médicaux à l'index FAISS
def add_medical_reference(file_content: bytes):
global index, documents
paragraphs = extract_text_from_pdf(file_content)
embeddings = model.encode(paragraphs)
index.add(np.array(embeddings, dtype="float32"))
documents.extend(paragraphs)
faiss.write_index(index, INDEX_PATH)
# Recherche dans FAISS pour trouver les documents pertinents
def search_faiss(query, k=5):
query_embedding = model.encode([query])
distances, indices = index.search(np.array(query_embedding, dtype="float32"), k)
results = [documents[i] for i in indices[0] if i < len(documents)]
return results
# Appel à l'API Gemini
def call_gemini_api(prompt):
headers = {"Content-Type": "application/json"}
payload = {
"contents": [
{
"parts": [
{"text": prompt}
]
}
]
}
try:
response = requests.post(f"{GEMINI_API_URL}?key={GEMINI_API_KEY}", json=payload, headers=headers)
response.raise_for_status()
response_json = response.json()
candidates = response_json.get("candidates", [])
if candidates:
return candidates[0]["content"]["parts"][0]["text"]
return "Pas de réponse disponible depuis Gemini."
except requests.exceptions.RequestException as e:
return f"Erreur API Gemini : {str(e)}"
# Ajouter un PDF de référence médicale
def upload_reference(file):
try:
# Ouvrir le fichier en mode binaire
with open(file.name, "rb") as f:
file_content = f.read()
add_medical_reference(file_content)
return "Référence médicale ajoutée avec succès."
except Exception as e:
return f"Erreur : {str(e)}"
# Analyser un PDF d'analyse de sang
def analyze_blood_test(file):
try:
# Ouvrir le fichier en mode binaire
with open(file.name, "rb") as f:
file_content = f.read()
# Extraire le texte des résultats d'analyse
test_results = extract_text_from_pdf(file_content)
if not test_results:
return "Aucun texte valide extrait du fichier d'analyse."
# Combine tous les résultats d'analyse dans un seul texte
query = "\n".join(test_results)
# Recherche dans l'index FAISS
relevant_docs = search_faiss(query, k=5)
context = "\n".join(relevant_docs)
# Enrichir le prompt avec les informations pertinentes
enriched_prompt = f"Voici les résultats d'analyse :\n{query}\n\nContexte pertinent :\n{context}"
gemini_response = call_gemini_api(enriched_prompt)
return {
"Réponse générée": gemini_response,
"Documents pertinents": relevant_docs
}
except Exception as e:
return f"Erreur : {str(e)}"
# Interface Gradio
def gradio_upload_reference(file):
return upload_reference(file)
def gradio_analyze_blood_test(file):
response = analyze_blood_test(file)
if isinstance(response, dict):
return f"Réponse générée :\n{response['Réponse générée']}\n\nDocuments pertinents :\n{response['Documents pertinents']}"
return response
# Lancer l'application Gradio
with gr.Blocks() as demo:
gr.Markdown("## Analyse Médicale avec RAG et Gemini")
with gr.Tab("Ajouter Références Médicales"):
ref_file = gr.File(label="Téléchargez un fichier PDF de référence médicale")
ref_output = gr.Textbox(label="Résultat")
ref_button = gr.Button("Ajouter")
ref_button.click(gradio_upload_reference, inputs=ref_file, outputs=ref_output)
with gr.Tab("Analyser un Résultat d'Analyse"):
test_file = gr.File(label="Téléchargez un fichier PDF d'analyse de sang")
analysis_output = gr.Textbox(label="Résultat")
analyze_button = gr.Button("Analyser")
analyze_button.click(gradio_analyze_blood_test, inputs=test_file, outputs=analysis_output)
demo.launch()