MahfoudAi's picture
fisrt version of RAG chatbot
c50b839
Raw
History Blame Contribute Delete
2.07 kB
import os
from langchain_text_splitters import MarkdownHeaderTextSplitter, RecursiveCharacterTextSplitter
from src.database import get_embeddings
from langchain_postgres import PGVector
from dotenv import load_dotenv
load_dotenv()
def load_and_split_data(file_path='data/reglement.md'):
"""Lit le fichier Markdown et le découpe sémantiquement en chunks."""
print(f"Lecture du fichier : {file_path}")
with open(file_path, "r", encoding="utf-8") as f:
markdown_document = f.read()
headers_to_split_on = [
("#", "Sujet_Principal"),
("##", "Type_Regle"),
]
markdown_splitter = MarkdownHeaderTextSplitter(headers_to_split_on=headers_to_split_on)
md_header_splits = markdown_splitter.split_text(markdown_document)
text_splitter = RecursiveCharacterTextSplitter(
chunk_size=700,
chunk_overlap=200,
separators=["\n\n", "\n", ".", " ", ""]
)
final_chunks = text_splitter.split_documents(md_header_splits)
print(f"Découpage terminé : {len(final_chunks)} chunks générés.")
return final_chunks
def build_index(chunks, db_url: str, collection_name: str):
"""Transforme les chunks en vecteurs avec bge-m3 et les stocke dans PostgreSQL."""
print("Chargement du modèle BAAI/bge-m3...")
embeddings = get_embeddings()
print(f"Connexion à la base et création de l'index vectoriel...")
vector_store = PGVector.from_documents(
embedding=embeddings,
documents=chunks,
collection_name=collection_name,
connection=db_url,
use_jsonb=True,
)
print("Indexation réussie dans la base de données PostgreSQL !")
return vector_store
if __name__ == "__main__":
FILE_PATH = "data/reglement.md"
COLLECTION_NAME = "rh_cgi_vectors"
DATABASE_URL = os.getenv('DATABASE_URL')
if not DATABASE_URL:
raise ValueError("Erreur : La variable DATABASE_URL est introuvable.")
documents = load_and_split_data(FILE_PATH)
vector_store = build_index(documents, DATABASE_URL, COLLECTION_NAME)