Skip to content

perf(faiss): batch Vector CSV embeddings #2759

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
greglum wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

perf(faiss): batch Vector CSV embeddings #2759

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
224 changes: 139 additions & 85 deletions py/vector_database/faiss/faiss_client.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -891,105 +891,83 @@ def _vector_addDocument(
"createdDocuments": [],
}

# loop through and embed new docs
# Embed each input CSV in one call, then preserve the existing per-Source
# dataset/vector artifacts used by removal and master-index rebuilding.
for document in documentFileLocation:
# Create the DataFrame for every file
dataset = self._load_dataset(dataset_location=document)

# Change the unit of work
# From being the document
# To the individual source inside the document
sources = dataset["Source"].unique().tolist()
directory = os.path.dirname(document)
effective_columns = (
list(dataset.columns)
if columns_to_index is None or len(columns_to_index) == 0
else columns_to_index
)
keyword_params = dict(keyword_search_params or {})
keyword_search = keyword_params.pop("keywordSearch", None) is True
prepared_sources = []

for source_name in sources:
source_dataset = dataset[dataset["Source"] == source_name].reset_index(
drop=True
)

# Get the directory path and the base filename without extension
directory, base_filename = os.path.split(document)
dataset_extension = ".parquet"
vector_extension = ".npy"

if columns_to_index == None or len(columns_to_index) == 0:
columns_to_index = list(source_dataset.columns)

# save the dataset, this is for efficiency after removing docs
new_file_path = os.path.join(
directory, source_name + "_dataset" + dataset_extension
if len(source_dataset) == 0:
continue
parts = source_dataset[effective_columns].astype(str)
source_dataset[target_column] = parts.apply(
lambda row: separator.join(row) + separator, axis=1
)

# if applicable, create the concatenated columns
if len(source_dataset) > 0:
# vectorized equivalent of mapping `_concatenate_columns`: build a
# target_column whose value for each row is
# `<v0><sep><v1><sep>...<vn-1><sep>` (note trailing separator,
# preserved for parity with the previous Dataset.map behavior).
parts = source_dataset[columns_to_index].astype(str)
source_dataset[target_column] = parts.apply(
lambda row: separator.join(row) + separator, axis=1
if keyword_search:
source_dataset[target_column] = (
self.keyword_engine.keyword_extraction(
input=list(source_dataset[target_column]),
insight_id=insight_id,
param_dict=keyword_params,
)
)

# transform chunks into keywords
if (
keyword_search_params != None
and keyword_search_params.pop("keywordSearch", None) is True
):
keywords_for_target_col = (
self.keyword_engine.keyword_extraction(
input=list(source_dataset[target_column]),
insight_id=insight_id,
param_dict=keyword_search_params,
)
vectors = self._embed_and_validate(
list(source_dataset[target_column]), insight_id
)
createDocumentsResponse["createdDocuments"].extend(
self._persist_source_artifacts(
directory,
source_name,
source_dataset,
vectors,
columns_to_remove,
target_column,
)
source_dataset[target_column] = keywords_for_target_col

# get the embeddings for the document
vectors = self.embeddings_engine.embeddings(
strings_to_embed=list(source_dataset[target_column]),
insight_id=insight_id,
)
vectors = np.array(vectors[0]["response"], dtype=np.float32)
assert vectors.ndim == 2
self._append_vectors(vectors)
else:
prepared_sources.append((source_name, source_dataset))

columns_to_remove.append(target_column)
columns_to_drop = list(
set(columns_to_remove).intersection(set(source_dataset.columns))
if prepared_sources:
all_text = [
value
for _source_name, source_dataset in prepared_sources
for value in source_dataset[target_column].tolist()
]
all_vectors = self._embed_and_validate(all_text, insight_id)
offset = 0
for source_name, source_dataset in prepared_sources:
source_rows = len(source_dataset)
source_vectors = all_vectors[offset : offset + source_rows]
offset += source_rows
createDocumentsResponse["createdDocuments"].extend(
self._persist_source_artifacts(
directory,
source_name,
source_dataset,
source_vectors,
columns_to_remove,
target_column,
)
)
source_dataset = source_dataset.drop(columns=columns_to_drop)

# Keep per-source files dtype-consistent so subsequent
# `_validateEmbeddingFiles` concatenations don't produce
# mixed-type columns that fail pyarrow serialization.
self._enforce_canonical_schema(source_dataset)
source_dataset.to_parquet(new_file_path, index=False)

# add the created source_dataset file path
createDocumentsResponse["createdDocuments"].append(new_file_path)

# normalize the vectors if using huggingface
if type(self.tokenizer).__name__ == "HuggingfaceTokenizer":
faiss.normalize_L2(vectors)

# write out the vectors as a .npy file (no pickle)
new_file_path = os.path.join(
directory,
source_name + "_vectors" + vector_extension,
if offset != len(all_vectors):
raise ValueError(
"Embedding rows could not be assigned to their Sources"
)
np.save(new_file_path, vectors, allow_pickle=False)

# add the created embeddings file path
createDocumentsResponse["createdDocuments"].append(new_file_path)

# TODO need to update the flow for how we instatiate
if self.encoded_vectors is None:
self.encoded_vectors = np.copy(vectors)
self.vector_dimensions = self.encoded_vectors.shape
else:
# make sure the dimensions are the same
assert self.vector_dimensions[1] == vectors.shape[1]
self.encoded_vectors = np.concatenate(
[self.encoded_vectors, vectors], axis=0
)
self._append_vectors(all_vectors)

master_indexClass_files, corrupted_file_sets = self.createMasterFiles(
path_to_files=os.path.dirname(os.path.dirname(documentFileLocation[0]))
Expand All @@ -1003,6 +981,82 @@ def _vector_addDocument(

return createDocumentsResponse

def _embed_and_validate(
self, strings_to_embed: List[str], insight_id: Optional[str]
) -> np.ndarray:
"""Embed one bounded CSV batch and validate it before writing artifacts."""
response = self.embeddings_engine.embeddings(
strings_to_embed=strings_to_embed,
insight_id=insight_id,
)
try:
vectors = np.asarray(response[0]["response"], dtype=np.float32)
except (IndexError, KeyError, TypeError, ValueError) as error:
raise ValueError(
"Embedding response does not contain a numeric response matrix"
) from error
if (
vectors.ndim != 2
or vectors.shape[0] != len(strings_to_embed)
or vectors.shape[1] == 0
):
raise ValueError(
"Embedding response shape does not match the submitted rows: "
f"expected {len(strings_to_embed)} rows, received {vectors.shape}"
)
if (
self.vector_dimensions is not None
and self.vector_dimensions[1] != vectors.shape[1]
):
raise ValueError(
"Embedding dimensions do not match the existing FAISS index: "
f"expected {self.vector_dimensions[1]}, received {vectors.shape[1]}"
)
if type(self.tokenizer).__name__ == "HuggingfaceTokenizer":
faiss.normalize_L2(vectors)
return vectors

def _persist_source_artifacts(
self,
directory: str,
source_name: str,
source_dataset: pd.DataFrame,
vectors: np.ndarray,
columns_to_remove: Optional[List[str]],
target_column: str,
) -> List[str]:
"""Write the stable dataset and vector pair for one Source."""
if len(source_dataset) != len(vectors):
raise ValueError(
f"Source {source_name} has {len(source_dataset)} rows but {len(vectors)} vectors"
)
columns_to_drop = list(
set([*(columns_to_remove or []), target_column]).intersection(
set(source_dataset.columns)
)
)
stored_dataset = source_dataset.drop(columns=columns_to_drop)
self._enforce_canonical_schema(stored_dataset)
dataset_path = os.path.join(directory, source_name + "_dataset.parquet")
vector_path = os.path.join(directory, source_name + "_vectors.npy")
stored_dataset.to_parquet(dataset_path, index=False)
np.save(vector_path, vectors, allow_pickle=False)
return [dataset_path, vector_path]

def _append_vectors(self, vectors: np.ndarray) -> None:
"""Maintain the in-memory vector shape until master files are rebuilt."""
if self.encoded_vectors is None:
self.encoded_vectors = np.copy(vectors)
self.vector_dimensions = self.encoded_vectors.shape
return
if self.vector_dimensions[1] != vectors.shape[1]:
raise ValueError(
"Embedding dimensions do not match the existing FAISS index: "
f"expected {self.vector_dimensions[1]}, received {vectors.shape[1]}"
)
self.encoded_vectors = np.concatenate([self.encoded_vectors, vectors], axis=0)
self.vector_dimensions = self.encoded_vectors.shape

def createMasterFiles(self, path_to_files: str) -> Tuple[str]:
"""
Create a master dataset and embeddings file based on the current documents. The main purpose of this is to improve startup runtime.
Expand Down
Loading
Loading