wehub-resource-sync
70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
from config import vector_collection, voyage_client, VOYAGE_MODEL
|
|
from pymongo.operations import SearchIndexModel
|
|
from langchain_community.document_loaders import PyPDFLoader
|
|
from langchain_text_splitters import RecursiveCharacterTextSplitter
|
|
import time
|
|
|
|
def get_embedding(data, input_type = "document"):
|
|
embeddings = voyage_client.embed(
|
|
data, model = VOYAGE_MODEL, input_type = input_type
|
|
).embeddings
|
|
return embeddings[0]
|
|
|
|
def ingest_data():
|
|
loader = PyPDFLoader("https://investors.mongodb.com/node/13176/pdf")
|
|
data = loader.load()
|
|
text_splitter = RecursiveCharacterTextSplitter(chunk_size=400, chunk_overlap=20)
|
|
documents = text_splitter.split_documents(data)
|
|
print(f"Successfully split PDF into {len(documents)} chunks.")
|
|
|
|
print("Generating embeddings and ingesting documents...")
|
|
docs_to_insert = []
|
|
for i, doc in enumerate(documents):
|
|
embedding = get_embedding(doc.page_content)
|
|
if embedding:
|
|
docs_to_insert.append({
|
|
"text": doc.page_content,
|
|
"embedding": embedding
|
|
})
|
|
|
|
if docs_to_insert:
|
|
result = vector_collection.insert_many(docs_to_insert)
|
|
print(f"Inserted {len(result.inserted_ids)} documents into the collection.")
|
|
else:
|
|
print("No documents were inserted. Check embedding generation process.")
|
|
|
|
index_name = "vector_index"
|
|
|
|
search_index_model = SearchIndexModel(
|
|
definition = {
|
|
"fields": [
|
|
{
|
|
"type": "vector",
|
|
"numDimensions": 1024,
|
|
"path": "embedding",
|
|
"similarity": "cosine"
|
|
}
|
|
]
|
|
},
|
|
name=index_name,
|
|
type="vectorSearch"
|
|
)
|
|
try:
|
|
vector_collection.create_search_index(model=search_index_model)
|
|
print(f"Search index '{index_name}' creation initiated.")
|
|
except Exception as e:
|
|
print(f"Error creating search index: {e}")
|
|
return
|
|
|
|
print("Polling to check if the index is ready. This may take up to a minute.")
|
|
predicate=None
|
|
if predicate is None:
|
|
predicate = lambda index: index.get("queryable") is True
|
|
|
|
while True:
|
|
indices = list(vector_collection.list_search_indexes(index_name))
|
|
if len(indices) and predicate(indices[0]):
|
|
break
|
|
time.sleep(5)
|
|
print(index_name + " is ready for querying.")
|