mirror of
https://github.com/imartinez/privateGPT.git
synced 2025-09-08 02:30:02 +00:00
Update code to use sentence-transformers through huggingfaceembeddings
This commit is contained in:
15
ingest.py
15
ingest.py
@@ -6,7 +6,7 @@ from dotenv import load_dotenv
|
||||
from langchain.document_loaders import TextLoader, PDFMinerLoader, CSVLoader
|
||||
from langchain.text_splitter import RecursiveCharacterTextSplitter
|
||||
from langchain.vectorstores import Chroma
|
||||
from langchain.embeddings import LlamaCppEmbeddings
|
||||
from langchain.embeddings import HuggingFaceEmbeddings
|
||||
from langchain.docstore.document import Document
|
||||
from constants import CHROMA_SETTINGS
|
||||
|
||||
@@ -38,22 +38,23 @@ def main():
|
||||
# Load environment variables
|
||||
persist_directory = os.environ.get('PERSIST_DIRECTORY')
|
||||
source_directory = os.environ.get('SOURCE_DIRECTORY', 'source_documents')
|
||||
llama_embeddings_model = os.environ.get('LLAMA_EMBEDDINGS_MODEL')
|
||||
model_n_ctx = os.environ.get('MODEL_N_CTX')
|
||||
embeddings_model_name = os.environ.get('EMBEDDINGS_MODEL_NAME')
|
||||
|
||||
# Load documents and split in chunks
|
||||
print(f"Loading documents from {source_directory}")
|
||||
chunk_size = 500
|
||||
chunk_overlap = 50
|
||||
documents = load_documents(source_directory)
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=500, chunk_overlap=50)
|
||||
text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
|
||||
texts = text_splitter.split_documents(documents)
|
||||
print(f"Loaded {len(documents)} documents from {source_directory}")
|
||||
print(f"Split into {len(texts)} chunks of text (max. 500 tokens each)")
|
||||
print(f"Split into {len(texts)} chunks of text (max. {chunk_size} characters each)")
|
||||
|
||||
# Create embeddings
|
||||
llama = LlamaCppEmbeddings(model_path=llama_embeddings_model, n_ctx=model_n_ctx)
|
||||
embeddings = HuggingFaceEmbeddings(model_name=embeddings_model_name)
|
||||
|
||||
# Create and store locally vectorstore
|
||||
db = Chroma.from_documents(texts, llama, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
|
||||
db = Chroma.from_documents(texts, embeddings, persist_directory=persist_directory, client_settings=CHROMA_SETTINGS)
|
||||
db.persist()
|
||||
db = None
|
||||
|
||||
|
Reference in New Issue
Block a user