Spaces:
Sleeping
Sleeping
File size: 1,897 Bytes
7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 e742fa6 7ffe358 |
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 |
"""
Modulo para procesar el PDF de la resolucion e indexar su contenido en la DB, para su posterior utilización por parte del chatbot.
Por simplicidad, se indexo un documento por cada página completa del documento. TODO: Implementar estrategia ParentDocumentRetriever.
"""
from langchain_community.document_loaders import PyPDFLoader
from chatbot.embeddings import init_embeddings
from chatbot.vectorstore import ChromaDB
from typing import List
import sys, argparse
def query(consulta: str) -> List:
""" Realiza una consulta a la vectorstore. """
embedding_model = init_embeddings()
vector_store = ChromaDB(embedding_model)
docs = vector_store.db.similarity_search(
consulta,
k=2,
)
return docs
def ingest(filename: str) -> None:
""" Ingiere un archivo PDF y lo indexa en la vectorstore. """
loader = PyPDFLoader(filename)
embedding_model = init_embeddings()
vector_store = ChromaDB(embedding_model)
for page in loader.lazy_load():
print(f"Procesando pagina {page.metadata['page']} - len: {len(page.page_content)}")
vector_store.add_documents([page])
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Este script permite interactuar con la vectorstore.')
parser.add_argument('operacion',
type=str,
choices=["ingest", "query"],
help="La operación que se desea realizar sobre la DB. 'ingest' para agregar un PDF a la DB. 'query' para hacer una consulta.")
parser.add_argument("argumento", help="Nombre del archivo a ingerir o consulta a realizar.")
args = parser.parse_args()
if args.operacion == 'ingest':
ingest(args.argumento)
print("Listo.")
elif args.operacion == 'query':
docs = query(args.argumento)
print(docs)
sys.exit(0)
|