feat: refactoring of rag node

This commit is contained in:
Marco Vinciguerra 2024-06-04 12:01:21 +02:00
parent dde0c7e27d
commit 7a13a6819f
2 changed files with 18 additions and 2 deletions

4
.gitignore vendored
View File

@ -23,6 +23,7 @@ docs/source/_static/
venv/
.venv/
.vscode/
.conda/
# exclude pdf, mp3
*.pdf
@ -38,3 +39,6 @@ lib/
*.html
.idea
# extras
cache/
run_smart_scraper.py

View File

@ -3,6 +3,7 @@ RAGNode Module
"""
from typing import List, Optional
import os
from langchain.docstore.document import Document
from langchain.retrievers import ContextualCompressionRetriever
@ -98,7 +99,18 @@ class RAGNode(BaseNode):
)
embeddings = self.embedder_model
retriever = FAISS.from_documents(chunked_docs, embeddings).as_retriever()
#------
index = FAISS.from_documents(chunked_docs, embeddings)
# Define the folder name
folder_name = "cache"
# Check if the folder exists, if not, create it
if not os.path.exists(folder_name):
os.makedirs(folder_name)
# Save the index to the folder
index.save_local(folder_name)
retriever = index.as_retriever()
#------
redundant_filter = EmbeddingsRedundantFilter(embeddings=embeddings)
# similarity_threshold could be set, now k=20
@ -121,4 +133,4 @@ class RAGNode(BaseNode):
self.logger.info("--- (tokens compressed and vector stored) ---")
state.update({self.output[0]: compressed_docs})
return state
return state