File size: 2,892 Bytes
92343a7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import os
import time
from urllib.request import Request, urlopen
from urllib.error import HTTPError, URLError
from pymongo import MongoClient
from langchain_community.document_loaders import PyPDFDirectoryLoader
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langchain_community.embeddings import HuggingFaceEmbeddings
from config import MONGODB_URI, DB_NAME, SAVE_FOLDER
PDF_URLS = [
]
COLLECTION_NAME = "connaissances"
def download_pdf(url, save_path, retries=2, delay=3):
"""Télécharge un PDF depuis une URL avec gestion des erreurs."""
for attempt in range(retries):
try:
req = Request(url, headers={'User-Agent': 'Mozilla/5.0'})
with urlopen(req) as response, open(save_path, 'wb') as f:
f.write(response.read())
print(f"Téléchargé : {save_path}")
return
except (HTTPError, URLError) as e:
print(f"Erreur ({e}) pour {url}, tentative {attempt+1}/{retries}")
time.sleep(delay)
print(f"Échec du téléchargement : {url}")
'''
def init_documents():
"""Initialise les documents dans la base de données avec leurs embeddings."""
os.makedirs(SAVE_FOLDER, exist_ok=True)
for url in PDF_URLS:
file_name = url.split("/")[-1]
file_path = os.path.join(SAVE_FOLDER, file_name)
if not os.path.exists(file_path):
download_pdf(url, file_path)
print("Chargement des PDFs...")
loader = PyPDFDirectoryLoader(SAVE_FOLDER)
docs = loader.load()
print("Découpage des documents...")
splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=100)
chunks = splitter.split_documents(docs)
print(f"{len(chunks)} morceaux extraits.")
print("Initialisation du modèle d'embeddings...")
embedding_model = HuggingFaceEmbeddings(model_name="shtilev/medical_embedded_v2")
print("Connexion à MongoDB...")
client = MongoClient(MONGODB_URI)
collection = client[DB_NAME][COLLECTION_NAME]
confirm = input("Cette opération supprimera toutes les données existantes. Continuer? (o/n): ")
if confirm.lower() != 'o':
print("Opération annulée.")
return
print("Suppression des documents existants...")
collection.delete_many({})
print("Génération des embeddings et insertion dans la base de données...")
for i, chunk in enumerate(chunks):
text = chunk.page_content
print(f"Traitement du morceau {i+1}/{len(chunks)}")
embedding = embedding_model.embed_query(text)
collection.insert_one({
"text": text,
"embedding": embedding
})
print("Tous les embeddings ont été insérés dans la base MongoDB.")
if __name__ == "__main__":
init_documents()
''' |