Spaces:
Sleeping
Sleeping
#----------- SETUP ----------- | |
from langchain_community.document_loaders import WebBaseLoader | |
from langchain_text_splitters import CharacterTextSplitter | |
from langchain_community.vectorstores import FAISS | |
from langchain_huggingface import HuggingFaceEmbeddings | |
from dotenv import load_dotenv | |
import os | |
import logging | |
logging.getLogger("langchain.text_splitter").setLevel(logging.ERROR) | |
import warnings | |
warnings.filterwarnings("ignore") | |
import yaml | |
# ----------- PARAMS ----------- | |
with open('./config.yaml', 'r', encoding='utf-8') as file: | |
config = yaml.safe_load(file) | |
EMBEDDING_MODEL = config.get('EMBEDDING_MODEL') | |
LLM_MODEL = config.get('LLM_MODEL') | |
REBUILD_VECTOR_STORE = config.get('REBUILD_VECTOR_STORE') | |
CHUNK_SIZE = config.get('CHUNK_SIZE') | |
CHUNK_OVERLAP = config.get('CHUNK_OVERLAP') | |
CACHE_FOLDER = config.get('CACHE_FOLDER') | |
URL_LIST = config.get('URL_LIST') | |
VS_BASE = config.get('VS_BASE') | |
# ----------- VECTOR STORE CREATION ----------- | |
def fn_rebuild_vector_store(REBUILD_VECTOR_STORE, URL_LIST, VS_BASE, EMBEDDING_MODEL, CACHE_FOLDER, CHUNK_SIZE, CHUNK_OVERLAP): | |
if REBUILD_VECTOR_STORE: | |
print("[INFO] REBUILD_VECTOR_STORE was set True. Recreating the vector store...") | |
loader = WebBaseLoader(web_paths=URL_LIST) | |
docs = loader.load() | |
text_splitter = CharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP) | |
split_docs = text_splitter.split_documents(docs) | |
embeddings = HuggingFaceEmbeddings( | |
model_name=EMBEDDING_MODEL, | |
cache_folder=CACHE_FOLDER) | |
vector_store = FAISS.from_documents(split_docs, embeddings) | |
os.makedirs(VS_BASE, exist_ok=True) | |
vector_store.save_local(VS_BASE) | |
print(f"[INFO] Vector Store saved in the path: {VS_BASE}") | |
else: | |
print("[INFO] REBUILD_VECTOR_STORE was set False. Using the current vector store...") | |
return print(f"[INFO] End of vector store process") | |