Docfile commited on
Commit
abe8f23
·
verified ·
1 Parent(s): f8c5a4c

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +89 -0
app.py ADDED
@@ -0,0 +1,89 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import streamlit as st
3
+ from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
4
+ from llama_index.llms.gemini import Gemini
5
+ from llama_index.core.extractors import TitleExtractor
6
+ from llama_index.core.node_parser import SentenceWindowNodeParser
7
+ from llama_index.core.ingestion import IngestionPipeline
8
+ from llama_index.core.query_engine import RetrieverQueryEngine
9
+ from llama_index.core.retrievers import AutoMergingRetriever
10
+ from llama_index.core.indices.vector_store.retrievers import VectorIndexRetriever
11
+ from llama_index.vector_stores.chroma import ChromaVectorStore
12
+
13
+ import chromadb
14
+
15
+ from dotenv import load_dotenv
16
+
17
+ load_dotenv()
18
+
19
+ # Configurer les paramètres globaux
20
+ Settings.llm = Gemini(api_key=os.environ["GOOGLE_API_KEY"], temperature=0.1)
21
+ Settings.chunk_size = 1024 # Taille des chunks pour l'indexation
22
+ # Nombre de tokens générés par le LLM
23
+
24
+ # Fonction pour charger les données et créer l'index (optimisé pour éviter les rechargements inutiles)
25
+ @st.cache_resource
26
+ def load_data_and_create_index():
27
+ """Charge les documents PDF et crée l'index vectoriel."""
28
+ documents = SimpleDirectoryReader("./data").load_data()
29
+
30
+ # Créer un pipeline d'ingestion avec extraction de titre et fenêtrage de phrases
31
+ node_parser = SentenceWindowNodeParser.from_defaults(
32
+ window_size=3,
33
+ window_metadata_key="window",
34
+ original_text_metadata_key="original_text",
35
+ )
36
+
37
+ text_splitter = node_parser.get_leaf_nodes_and_parent_nodes
38
+ extractors = [TitleExtractor(nodes=5)]
39
+
40
+ pipeline = IngestionPipeline(
41
+ transformations=[node_parser, *extractors]
42
+ )
43
+
44
+ # Indexer les documents
45
+ nodes = pipeline.run(documents=documents)
46
+
47
+ # Initialiser la base de données vectorielle (exemple avec Chroma)
48
+ db = chromadb.Client()
49
+ chroma_collection = db.get_or_create_collection("legal_docs")
50
+ vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
51
+
52
+ # Créer l'index
53
+ index = VectorStoreIndex.from_documents(nodes, vector_store=vector_store)
54
+ return index
55
+
56
+ # Fonction pour effectuer la requête
57
+ def perform_query(query_str, index):
58
+ """Effectue une requête sur l'index et renvoie la réponse."""
59
+ # Créer un AutoMergingRetriever
60
+ base_retriever = VectorIndexRetriever(
61
+ index=index,
62
+ similarity_top_k=8,
63
+ )
64
+ retriever = AutoMergingRetriever(base_retriever, index.storage_context)
65
+
66
+ # Créer le moteur de requête
67
+ query_engine = RetrieverQueryEngine.from_args(retriever=retriever)
68
+
69
+ response = query_engine.query(query_str)
70
+ return response
71
+
72
+ # Interface utilisateur Streamlit
73
+ st.title("Agent de Questions-Réponses Juridiques (Gemini + LlamaIndex)")
74
+
75
+ # Charger les données et créer l'index (une seule fois)
76
+ index = load_data_and_create_index()
77
+
78
+ # Champ de saisie de la question
79
+ query_str = st.text_input("Posez votre question juridique ici :")
80
+
81
+ # Bouton pour soumettre la question
82
+ if st.button("Poser la question"):
83
+ if query_str:
84
+ with st.spinner("Recherche en cours..."):
85
+ response = perform_query(query_str, index)
86
+ st.success("Réponse :")
87
+ st.write(response)
88
+ else:
89
+ st.error("Veuillez saisir une question.")