|
|
|
import tables |
|
from pathlib import Path |
|
import numpy as np |
|
from numpy.lib.recfunctions import structured_to_unstructured |
|
from tqdm import tqdm |
|
from numba import njit, prange |
|
from time import time |
|
|
|
|
|
|
|
@njit('float32[:](uint8[:])', parallel=True) |
|
def tofp32n8(arr): |
|
"""Numba-optimized function that converts a fp8 (4M3E) array to fp32 using a mapping table |
|
The array is assumed to be one dimensional with the fp8 |
|
represented as UInt8 |
|
""" |
|
fp8table= np.frombuffer(b'\x00\x00\x00\x00\x00\x00\x00;\x00\x00\x80;\x00\x00\xc0;\x00\x00\x00<\x00\x00 <\x00\x00@<\x00\x00`<\x00\x00\x80<\x00\x00\x90<\x00\x00\xa0<\x00\x00\xb0<\x00\x00\xc0<\x00\x00\xd0<\x00\x00\xe0<\x00\x00\xf0<\x00\x00\x00=\x00\x00\x10=\x00\x00 =\x00\x000=\x00\x00@=\x00\x00P=\x00\x00`=\x00\x00p=\x00\x00\x80=\x00\x00\x90=\x00\x00\xa0=\x00\x00\xb0=\x00\x00\xc0=\x00\x00\xd0=\x00\x00\xe0=\x00\x00\xf0=\x00\x00\x00>\x00\x00\x10>\x00\x00 >\x00\x000>\x00\x00@>\x00\x00P>\x00\x00`>\x00\x00p>\x00\x00\x80>\x00\x00\x90>\x00\x00\xa0>\x00\x00\xb0>\x00\x00\xc0>\x00\x00\xd0>\x00\x00\xe0>\x00\x00\xf0>\x00\x00\x00?\x00\x00\x10?\x00\x00 ?\x00\x000?\x00\x00@?\x00\x00P?\x00\x00`?\x00\x00p?\x00\x00\x80?\x00\x00\x90?\x00\x00\xa0?\x00\x00\xb0?\x00\x00\xc0?\x00\x00\xd0?\x00\x00\xe0?\x00\x00\xf0?\x00\x00\x00@\x00\x00\x10@\x00\x00 @\x00\x000@\x00\x00@@\x00\x00P@\x00\x00`@\x00\x00p@\x00\x00\x80@\x00\x00\x90@\x00\x00\xa0@\x00\x00\xb0@\x00\x00\xc0@\x00\x00\xd0@\x00\x00\xe0@\x00\x00\xf0@\x00\x00\x00A\x00\x00\x10A\x00\x00 A\x00\x000A\x00\x00@A\x00\x00PA\x00\x00`A\x00\x00pA\x00\x00\x80A\x00\x00\x90A\x00\x00\xa0A\x00\x00\xb0A\x00\x00\xc0A\x00\x00\xd0A\x00\x00\xe0A\x00\x00\xf0A\x00\x00\x00B\x00\x00\x10B\x00\x00 B\x00\x000B\x00\x00@B\x00\x00PB\x00\x00`B\x00\x00pB\x00\x00\x80B\x00\x00\x90B\x00\x00\xa0B\x00\x00\xb0B\x00\x00\xc0B\x00\x00\xd0B\x00\x00\xe0B\x00\x00\xf0B\x00\x00\x00C\x00\x00\x10C\x00\x00 C\x00\x000C\x00\x00@C\x00\x00PC\x00\x00`C\x00\x00pC\x00\x00\x80C\x00\x00\x90C\x00\x00\xa0C\x00\x00\xb0C\x00\x00\xc0C\x00\x00\xd0C\x00\x00\xe0C\x00\x00\xf0C\x00\x00\x00\x80\x00\x00\x00\xbb\x00\x00\x80\xbb\x00\x00\xc0\xbb\x00\x00\x00\xbc\x00\x00 \xbc\x00\x00@\xbc\x00\x00`\xbc\x00\x00\x80\xbc\x00\x00\x90\xbc\x00\x00\xa0\xbc\x00\x00\xb0\xbc\x00\x00\xc0\xbc\x00\x00\xd0\xbc\x00\x00\xe0\xbc\x00\x00\xf0\xbc\x00\x00\x00\xbd\x00\x00\x10\xbd\x00\x00 \xbd\x00\x000\xbd\x00\x00@\xbd\x00\x00P\xbd\x00\x00`\xbd\x00\x00p\xbd\x00\x00\x80\xbd\x00\x00\x90\xbd\x00\x00\xa0\xbd\x00\x00\xb0\xbd\x00\x00\xc0\xbd\x00\x00\xd0\xbd\x00\x00\xe0\xbd\x00\x00\xf0\xbd\x00\x00\x00\xbe\x00\x00\x10\xbe\x00\x00 \xbe\x00\x000\xbe\x00\x00@\xbe\x00\x00P\xbe\x00\x00`\xbe\x00\x00p\xbe\x00\x00\x80\xbe\x00\x00\x90\xbe\x00\x00\xa0\xbe\x00\x00\xb0\xbe\x00\x00\xc0\xbe\x00\x00\xd0\xbe\x00\x00\xe0\xbe\x00\x00\xf0\xbe\x00\x00\x00\xbf\x00\x00\x10\xbf\x00\x00 \xbf\x00\x000\xbf\x00\x00@\xbf\x00\x00P\xbf\x00\x00`\xbf\x00\x00p\xbf\x00\x00\x80\xbf\x00\x00\x90\xbf\x00\x00\xa0\xbf\x00\x00\xb0\xbf\x00\x00\xc0\xbf\x00\x00\xd0\xbf\x00\x00\xe0\xbf\x00\x00\xf0\xbf\x00\x00\x00\xc0\x00\x00\x10\xc0\x00\x00 \xc0\x00\x000\xc0\x00\x00@\xc0\x00\x00P\xc0\x00\x00`\xc0\x00\x00p\xc0\x00\x00\x80\xc0\x00\x00\x90\xc0\x00\x00\xa0\xc0\x00\x00\xb0\xc0\x00\x00\xc0\xc0\x00\x00\xd0\xc0\x00\x00\xe0\xc0\x00\x00\xf0\xc0\x00\x00\x00\xc1\x00\x00\x10\xc1\x00\x00 \xc1\x00\x000\xc1\x00\x00@\xc1\x00\x00P\xc1\x00\x00`\xc1\x00\x00p\xc1\x00\x00\x80\xc1\x00\x00\x90\xc1\x00\x00\xa0\xc1\x00\x00\xb0\xc1\x00\x00\xc0\xc1\x00\x00\xd0\xc1\x00\x00\xe0\xc1\x00\x00\xf0\xc1\x00\x00\x00\xc2\x00\x00\x10\xc2\x00\x00 \xc2\x00\x000\xc2\x00\x00@\xc2\x00\x00P\xc2\x00\x00`\xc2\x00\x00p\xc2\x00\x00\x80\xc2\x00\x00\x90\xc2\x00\x00\xa0\xc2\x00\x00\xb0\xc2\x00\x00\xc0\xc2\x00\x00\xd0\xc2\x00\x00\xe0\xc2\x00\x00\xf0\xc2\x00\x00\x00\xc3\x00\x00\x10\xc3\x00\x00 \xc3\x00\x000\xc3\x00\x00@\xc3\x00\x00P\xc3\x00\x00`\xc3\x00\x00p\xc3\x00\x00\x80\xc3\x00\x00\x90\xc3\x00\x00\xa0\xc3\x00\x00\xb0\xc3\x00\x00\xc0\xc3\x00\x00\xd0\xc3\x00\x00\xe0\xc3\x00\x00\x00\x00\x00\x00\x00;\x00\x00\x80;\x00\x00\xc0;\x00\x00\x00<\x00\x00 <\x00\x00@<\x00\x00`<\x00\x00\x80<\x00\x00\x90<\x00\x00\xa0<\x00\x00\xb0<\x00\x00\xc0<\x00\x00\xd0<\x00\x00\xe0<\x00\x00\xf0<\x00\x00\x00=\x00\x00\x10=\x00\x00 =\x00\x000=\x00\x00@=\x00\x00P=\x00\x00`=\x00\x00p=\x00\x00\x80=\x00\x00\x90=\x00\x00\xa0=\x00\x00\xb0=\x00\x00\xc0=\x00\x00\xd0=\x00\x00\xe0=\x00\x00\xf0=\x00\x00\x00>\x00\x00\x10>\x00\x00 >\x00\x000>\x00\x00@>\x00\x00P>\x00\x00`>\x00\x00p>\x00\x00\x80>\x00\x00\x90>\x00\x00\xa0>\x00\x00\xb0>\x00\x00\xc0>\x00\x00\xd0>\x00\x00\xe0>\x00\x00\xf0>\x00\x00\x00?\x00\x00\x10?\x00\x00 ?\x00\x000?\x00\x00@?\x00\x00P?\x00\x00`?\x00\x00p?\x00\x00\x80?\x00\x00\x90?\x00\x00\xa0?\x00\x00\xb0?\x00\x00\xc0?\x00\x00\xd0?\x00\x00\xe0?\x00\x00\xf0?\x00\x00\x00@\x00\x00\x10@\x00\x00 @\x00\x000@\x00\x00@@\x00\x00P@\x00\x00`@\x00\x00p@\x00\x00\x80@\x00\x00\x90@\x00\x00\xa0@\x00\x00\xb0@\x00\x00\xc0@\x00\x00\xd0@\x00\x00\xe0@\x00\x00\xf0@\x00\x00\x00A\x00\x00\x10A\x00\x00 A\x00\x000A\x00\x00@A\x00\x00PA\x00\x00`A\x00\x00pA\x00\x00\x80A\x00\x00\x90A\x00\x00\xa0A\x00\x00\xb0A\x00\x00\xc0A\x00\x00\xd0A\x00\x00\xe0A\x00\x00\xf0A\x00\x00\x00B\x00\x00\x10B\x00\x00 B\x00\x000B\x00\x00@B\x00\x00PB\x00\x00`B\x00\x00pB\x00\x00\x80B\x00\x00\x90B\x00\x00\xa0B\x00\x00\xb0B\x00\x00\xc0B\x00\x00\xd0B\x00\x00\xe0B\x00\x00\xf0B\x00\x00\x00C\x00\x00\x10C\x00\x00 C\x00\x000C\x00\x00@C\x00\x00PC\x00\x00`C\x00\x00pC\x00\x00\x80C\x00\x00\x90C\x00\x00\xa0C\x00\x00\xb0C\x00\x00\xc0C\x00\x00\xd0C\x00\x00\xe0C\x00\x00\xf0C\x00\x00\x00\x80\x00\x00\x00\xbb\x00\x00\x80\xbb\x00\x00\xc0\xbb\x00\x00\x00\xbc\x00\x00 \xbc\x00\x00@\xbc\x00\x00`\xbc\x00\x00\x80\xbc\x00\x00\x90\xbc\x00\x00\xa0\xbc\x00\x00\xb0\xbc\x00\x00\xc0\xbc\x00\x00\xd0\xbc\x00\x00\xe0\xbc\x00\x00\xf0\xbc\x00\x00\x00\xbd\x00\x00\x10\xbd\x00\x00 \xbd\x00\x000\xbd\x00\x00@\xbd\x00\x00P\xbd\x00\x00`\xbd\x00\x00p\xbd\x00\x00\x80\xbd\x00\x00\x90\xbd\x00\x00\xa0\xbd\x00\x00\xb0\xbd\x00\x00\xc0\xbd\x00\x00\xd0\xbd\x00\x00\xe0\xbd\x00\x00\xf0\xbd\x00\x00\x00\xbe\x00\x00\x10\xbe\x00\x00 \xbe\x00\x000\xbe\x00\x00@\xbe\x00\x00P\xbe\x00\x00`\xbe\x00\x00p\xbe\x00\x00\x80\xbe\x00\x00\x90\xbe\x00\x00\xa0\xbe\x00\x00\xb0\xbe\x00\x00\xc0\xbe\x00\x00\xd0\xbe\x00\x00\xe0\xbe\x00\x00\xf0\xbe\x00\x00\x00\xbf\x00\x00\x10\xbf\x00\x00 \xbf\x00\x000\xbf\x00\x00@\xbf\x00\x00P\xbf\x00\x00`\xbf\x00\x00p\xbf\x00\x00\x80\xbf\x00\x00\x90\xbf\x00\x00\xa0\xbf\x00\x00\xb0\xbf\x00\x00\xc0\xbf\x00\x00\xd0\xbf\x00\x00\xe0\xbf\x00\x00\xf0\xbf\x00\x00\x00\xc0\x00\x00\x10\xc0\x00\x00 \xc0\x00\x000\xc0\x00\x00@\xc0\x00\x00P\xc0\x00\x00`\xc0\x00\x00p\xc0\x00\x00\x80\xc0\x00\x00\x90\xc0\x00\x00\xa0\xc0\x00\x00\xb0\xc0\x00\x00\xc0\xc0\x00\x00\xd0\xc0\x00\x00\xe0\xc0\x00\x00\xf0\xc0\x00\x00\x00\xc1\x00\x00\x10\xc1\x00\x00 \xc1\x00\x000\xc1\x00\x00@\xc1\x00\x00P\xc1\x00\x00`\xc1\x00\x00p\xc1\x00\x00\x80\xc1\x00\x00\x90\xc1\x00\x00\xa0\xc1\x00\x00\xb0\xc1\x00\x00\xc0\xc1\x00\x00\xd0\xc1\x00\x00\xe0\xc1\x00\x00\xf0\xc1\x00\x00\x00\xc2\x00\x00\x10\xc2\x00\x00 \xc2\x00\x000\xc2\x00\x00@\xc2\x00\x00P\xc2\x00\x00`\xc2\x00\x00p\xc2\x00\x00\x80\xc2\x00\x00\x90\xc2\x00\x00\xa0\xc2\x00\x00\xb0\xc2\x00\x00\xc0\xc2\x00\x00\xd0\xc2\x00\x00\xe0\xc2\x00\x00\xf0\xc2\x00\x00\x00\xc3\x00\x00\x10\xc3\x00\x00 \xc3\x00\x000\xc3\x00\x00@\xc3\x00\x00P\xc3\x00\x00`\xc3\x00\x00p\xc3\x00\x00\x80\xc3\x00\x00\x90\xc3\x00\x00\xa0\xc3\x00\x00\xb0\xc3\x00\x00\xc0\xc3\x00\x00\xd0\xc3\x00\x00\xe0\xc3', dtype=np.float32) |
|
arr2 = np.empty(arr.shape[0], dtype="float32") |
|
for i in prange(arr.shape[0]): |
|
arr2[i] = fp8table[arr[i]] |
|
return arr2 |
|
|
|
|
|
def tofp32(arr): |
|
"""Converts a fp8 (4M3E) array to fp32. |
|
Reshapes the array to be one |
|
dimensional and uses a numba-optimized function |
|
""" |
|
return tofp32n8(arr.reshape(arr.shape[0]*arr.shape[1])).reshape(arr.shape) |
|
|
|
|
|
@njit('uint8[:](uint32[:])', parallel=True) |
|
def tofp8n(arr): |
|
"""Numba-optimized function that converts an array of fp32 to fp8 (4M3E) |
|
Uses the algorithm described by ProjectPhysX at https://stackoverflow.com/questions/1659440/32-bit-to-16-bit-floating-point-conversion |
|
and https://www.researchgate.net/publication/362275548_Accuracy_and_performance_of_the_lattice_Boltzmann_method_with_64-bit_32-bit_and_customized_16-bit_number_formats |
|
""" |
|
arr2 = np.empty(arr.shape[0], dtype="uint8") |
|
for i in prange(arr.shape[0]): |
|
|
|
y = arr[i] + 0x00080000 |
|
e = (y&0x7F800000)>>23 |
|
m = y&0x007FFFFF |
|
|
|
if e > 135: |
|
arr2[i] = 0x7F | (y&0x80000000)>>24 |
|
elif e > 120: |
|
arr2[i] = ((e-120)<<3) & 0x78 | m>>20 | (y&0x80000000)>>24 |
|
elif e < 121 and e > 116: |
|
|
|
arr2[i] = ((((m+0x00780000)>>(140-e))+1)>>1) | (y&0x80000000)>>24 |
|
else: |
|
arr2[i] = 0 | (y&0x80000000)>>24 |
|
return arr2 |
|
|
|
|
|
def tofp8(arr): |
|
"""Converts an array of fp32 to fp8 (4M3E) |
|
Reshapes the array to be one |
|
dimensional and uses a numba-optimized function |
|
""" |
|
return tofp8n(arr.view(dtype=np.uint32).reshape(arr.shape[0]*arr.shape[1])).view(dtype=np.uint8).reshape(arr.shape) |
|
|
|
|
|
|
|
class BlobTable(): |
|
"""Class to handle a storage of variable-length values of a key-value storage |
|
Key is fixed length of key_length |
|
""" |
|
def __init__(self, store, key_length=20): |
|
"""Initializes class using a pytables store and a key_length value |
|
""" |
|
if "keys" not in store.root: |
|
|
|
filters = tables.Filters(complevel=5, complib='blosc:lz4', |
|
shuffle=1, bitshuffle=0) |
|
|
|
blob_type = {"key": tables.StringCol(key_length, pos=0), |
|
"offset":tables.Int64Col(pos=1), |
|
"length": tables.Int64Col(pos=2), |
|
} |
|
|
|
self.keys_table = store.create_table("/", "keys", |
|
blob_type, |
|
filters=filters, |
|
chunkshape=10000) |
|
self.values_table = store.create_earray("/", "values", atom=tables.UInt8Atom(), shape=(0,), filters=filters) |
|
else: |
|
self.keys_table = store.root.keys |
|
self.values_table = store.root.values |
|
|
|
self.offset = self.values_table.nrows |
|
self.nrows = self.keys_table.nrows |
|
self._is_closed = False |
|
|
|
|
|
def __len__(self): |
|
return self.nrows |
|
|
|
def create_index(self): |
|
self.keys_table.cols.key.reindex() |
|
|
|
|
|
def append(self, key, value): |
|
"""Appends a key-value to the storage |
|
""" |
|
|
|
length = len(value) |
|
self.values_table.append(np.frombuffer(value, dtype=np.uint8)) |
|
|
|
|
|
row = self.keys_table.row |
|
row["key"] = key |
|
row["offset"] = self.offset |
|
row["length"] = length |
|
row.append() |
|
self.offset += length |
|
self.nrows += 1 |
|
|
|
def __getitem__ (self, rownum): |
|
if isinstance(rownum, slice): |
|
return [self[ii] for ii in range(*rownum.indices(len(self)))] |
|
else: |
|
row = self.keys_table[rownum] |
|
offset = row['offset'] |
|
value = self.values_table.read(offset, offset+row["length"]).tobytes() |
|
|
|
return value |
|
|
|
|
|
def get_value (self, key): |
|
key = key.encode("utf8") |
|
offset, length = [(r['offset'], r['length']) for r in self.keys_table.where(f"key=={key}")][0] |
|
value = self.values_table.read(offset, offset+length).tobytes() |
|
return value |
|
|
|
|
|
class Simiandb(): |
|
"""Wrapper around pytables store . |
|
To use, you should have the ``pytables`` python package installed. |
|
Example: |
|
.. code-block:: python |
|
from simiandb import Simiandb |
|
docdb = simiandb("store") |
|
""" |
|
|
|
def __init__(self, storepath, embedding_function=None, mode="a", id_length = 19): |
|
|
|
if mode not in ["a", "w", "r"]: |
|
raise ValueError("Mode can only be r, w or a") |
|
self._embedding_function = embedding_function |
|
self._storename = Path(storepath) |
|
self._mode = mode |
|
if not self._storename.exists(): |
|
self._storename.mkdir() |
|
|
|
self._vectorstore = tables.open_file( self._storename / "embeddings.h5", mode = mode) |
|
self._docstore = tables.open_file( self._storename / "documents.h5", mode = mode) |
|
self._metastore = tables.open_file( self._storename / "metadatas.h5", mode = mode) |
|
self._embedding_function = embedding_function |
|
self._is_closed = False |
|
if 'embeddings' in self._vectorstore.root: |
|
self._vector_table = self._vectorstore.root.embeddings |
|
self._docs_table = BlobTable(self._docstore, id_length) |
|
return |
|
|
|
|
|
def __enter__(self): |
|
"""Magic method Required for usage with the with statement |
|
""" |
|
return self |
|
|
|
|
|
def _get_top_indexes(self, c, k): |
|
count = self._vector_table.nrows |
|
st =0 |
|
batch = self._vector_table.chunkshape[0]*25 |
|
res = np.ascontiguousarray(np.empty(shape=(count,), dtype="float32")) |
|
end = 0 |
|
|
|
while end!=count: |
|
end += batch |
|
end = end if end <= count else count |
|
t_res = structured_to_unstructured(self._vector_table.read(start=st, stop=end)) |
|
t_res = tofp32(t_res) |
|
np.dot(t_res,c, res[st:end]) |
|
st = end |
|
|
|
indices = np.argpartition(res, -k)[-k:] |
|
indices = indices[np.argsort(res[indices])[::-1]] |
|
|
|
return indices |
|
|
|
|
|
def _create_embeddings_table(self, dimensions): |
|
"""Creates the embeddings table within the pytables file |
|
""" |
|
if dimensions > 512: |
|
|
|
tables.parameters.MAX_COLUMNS = len(dimensions) |
|
embedding_type = {f"d{n}":tables.UInt8Col(pos=n) for n in range(dimensions)} |
|
|
|
|
|
filters = None |
|
|
|
self._vector_table = self._vectorstore.create_table("/", "embeddings", |
|
embedding_type, |
|
filters=filters, |
|
chunkshape=10000) |
|
|
|
|
|
def _check_closed(self): |
|
if self._is_closed: |
|
raise ValueError("Simiandb is already closed") |
|
|
|
|
|
def add_texts(self, texts, metadatas = None, ids = None, embeddings=None, show_progressbar=True): |
|
"""Run more texts through the embeddings and add to the vectorstore. |
|
Args: |
|
texts (Iterable[str]): Texts to add to the vectorstore. |
|
metadatas (Optional[List[dict]], optional): Optional list of metadatas. |
|
ids (Optional[List[str]], optional): Optional list of IDs. |
|
embeddings (Optional[List[array]], optional): Optional list of embeddings. |
|
Returns: |
|
List[str]: List of IDs of the added texts. |
|
""" |
|
|
|
self._check_closed() |
|
|
|
self._add_embeddings(texts, embeddings, show_progressbar) |
|
|
|
if ids is None: |
|
ids = list(range(self.docs_table.nrows, self.docs_table.nrows + len(texts))) |
|
|
|
for textid, text in zip(ids, texts): |
|
self.docs_table.append(textid, text.encode("utf8")) |
|
|
|
return ids |
|
|
|
|
|
def get_text(self, key): |
|
return self._docs_table.get_value(key).decode("utf8") |
|
|
|
|
|
def create_keys_index(self): |
|
self._docs_table.create_index() |
|
|
|
|
|
def _add_embeddings(self, texts, embeddings, show_progressbar): |
|
"""Calculate or use embeddings to fill the embeddings table |
|
""" |
|
|
|
if embeddings is None and not self._embedding_function is None: |
|
embeddings = self._embedding_function.embed_documents(texts) |
|
|
|
if not embeddings is None and 'embeddings' not in self._vectorstore.root: |
|
dimensions = len(embeddings[0]) |
|
self._create_embeddings_table(dimensions) |
|
|
|
if not embeddings is None : |
|
self._vector_table = self._vectorstore.root.embeddings |
|
embeddings = tofp8(np.array(embeddings, dtype=np.float32)) |
|
self._vector_table.append(embeddings) |
|
|
|
|
|
def regenerate_embeddings(self, embeddings=None, show_progressbar=True): |
|
"""Run existing texts through the embeddings and add to the vectorstore. |
|
Args: |
|
embeddings (Optional[List[array]], optional): Optional list of embeddings. |
|
""" |
|
|
|
self._check_closed() |
|
self._vectorstore.close() |
|
(self._storename / "embeddings.h5").kill() |
|
self._vectorstore = tables.open_file( self._storename / "embeddings.h5", mode = self._mode) |
|
|
|
batch_size = 1000 |
|
for i in tqdm(range(0, len(self.docs_table), batch_size), disable=not show_progressbar): |
|
text_batch = [text.decode("utf8") for text in self.docs_table[i:i+batch_size]] |
|
if embeddings is not None: |
|
embeddings_batch = embeddings[i:i+batch_size] |
|
elif self.embedding_function is not None: |
|
embeddings_batch = self._embedding_function.embed_documents(text_batch) |
|
else: |
|
raise ValueError("Neither embeddings nor embedding function provided") |
|
self._add_embeddings(text_batch, embeddings_batch, show_progressbar) |
|
return |
|
|
|
|
|
def similarity_search(self, query: str, k = 4, filter = None): |
|
"""Run similarity search with PytableStore. |
|
Args: |
|
query (str): Query text to search for. |
|
k (int): Number of results to return. Defaults to 4. |
|
filter (Optional[Dict[str, str]]): Filter by metadata. Defaults to None. |
|
Returns: |
|
List[Document]: List of documents most simmilar to the query text. |
|
""" |
|
self._check_closed() |
|
query_embedding = np.array(self._embedding_function.embed_query(query),dtype="float32") |
|
results = self._get_top_indexes(query_embedding, k) |
|
|
|
docs = [self._docs_table[i].decode("utf8") for i in results] |
|
return docs |
|
|
|
|
|
def close(self): |
|
"""Makes sure the pytables file is closed |
|
""" |
|
if not self._is_closed: |
|
self._is_closed = True |
|
|
|
if hasattr(self, '_Simiandb__vectorstore'): |
|
try: |
|
self._vectorstore.flush() |
|
self._docstore.flush() |
|
self._metastore.flush() |
|
self._vectorstore.close() |
|
self._docstore.close() |
|
self._metastore.close() |
|
except: |
|
print("Unable to close file") |
|
|
|
|
|
def __exit__(self, exc_type, exc_value, exc_traceback): |
|
"""Magic method Required for usage with the with statement |
|
""" |
|
self.close() |
|
|
|
def __del__(self): |
|
"""Magic method just in case the object is deleted without closing it |
|
""" |
|
self.close() |
|
|
|
|
|
|
|
if __name__ == '__main__': |
|
pass |
|
|
