File size: 3,216 Bytes
abe8f23 |
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 82 83 84 85 86 87 88 89 |
import os
import streamlit as st
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.gemini import Gemini
from llama_index.core.extractors import TitleExtractor
from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core.indices.vector_store.retrievers import VectorIndexRetriever
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb
from dotenv import load_dotenv
load_dotenv()
# Configurer les paramètres globaux
Settings.llm = Gemini(api_key=os.environ["GOOGLE_API_KEY"], temperature=0.1)
Settings.chunk_size = 1024 # Taille des chunks pour l'indexation
# Nombre de tokens générés par le LLM
# Fonction pour charger les données et créer l'index (optimisé pour éviter les rechargements inutiles)
@st.cache_resource
def load_data_and_create_index():
"""Charge les documents PDF et crée l'index vectoriel."""
documents = SimpleDirectoryReader("./data").load_data()
# Créer un pipeline d'ingestion avec extraction de titre et fenêtrage de phrases
node_parser = SentenceWindowNodeParser.from_defaults(
window_size=3,
window_metadata_key="window",
original_text_metadata_key="original_text",
)
text_splitter = node_parser.get_leaf_nodes_and_parent_nodes
extractors = [TitleExtractor(nodes=5)]
pipeline = IngestionPipeline(
transformations=[node_parser, *extractors]
)
# Indexer les documents
nodes = pipeline.run(documents=documents)
# Initialiser la base de données vectorielle (exemple avec Chroma)
db = chromadb.Client()
chroma_collection = db.get_or_create_collection("legal_docs")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
# Créer l'index
index = VectorStoreIndex.from_documents(nodes, vector_store=vector_store)
return index
# Fonction pour effectuer la requête
def perform_query(query_str, index):
"""Effectue une requête sur l'index et renvoie la réponse."""
# Créer un AutoMergingRetriever
base_retriever = VectorIndexRetriever(
index=index,
similarity_top_k=8,
)
retriever = AutoMergingRetriever(base_retriever, index.storage_context)
# Créer le moteur de requête
query_engine = RetrieverQueryEngine.from_args(retriever=retriever)
response = query_engine.query(query_str)
return response
# Interface utilisateur Streamlit
st.title("Agent de Questions-Réponses Juridiques (Gemini + LlamaIndex)")
# Charger les données et créer l'index (une seule fois)
index = load_data_and_create_index()
# Champ de saisie de la question
query_str = st.text_input("Posez votre question juridique ici :")
# Bouton pour soumettre la question
if st.button("Poser la question"):
if query_str:
with st.spinner("Recherche en cours..."):
response = perform_query(query_str, index)
st.success("Réponse :")
st.write(response)
else:
st.error("Veuillez saisir une question.") |