def load_searcher(db, artifacts_dir):
tokenization_path = os.path.join(artifacts_dir, "datapoint_to_token.npy")
hashed_db_path = os.path.join(artifacts_dir, "hashed_dataset.npy")
tokenization = np.load(tokenization_path) if os.path.isfile(
tokenization_path) else None
hashed_db = np.load(hashed_db_path) if os.path.isfile(
hashed_db_path) else None
return ScannSearcher(
scann_pybind.ScannNumpy(db, tokenization, hashed_db, artifacts_dir))
def load_searcher(artifacts_dir):
"""Loads searcher assets from artifacts_dir and returns a ScaNN searcher."""
def load_if_exists(filename):
path = os.path.join(artifacts_dir, filename)
return np.load(path) if os.path.isfile(path) else None
db = load_if_exists("dataset.npy")
tokenization = load_if_exists("datapoint_to_token.npy")
hashed_db = load_if_exists("hashed_dataset.npy")
int8_db = load_if_exists("int8_dataset.npy")
int8_multipliers = load_if_exists("int8_multipliers.npy")
db_norms = load_if_exists("dp_norms.npy")
return ScannSearcher(
scann_pybind.ScannNumpy(db, tokenization, hashed_db, int8_db,
int8_multipliers, db_norms, artifacts_dir))
def create_searcher(db,
train_set,
scann_config,
load_coarse,
coarse_path,
load_fine,
fine_path,
training_threads=0):
# if load_coarse == True:
# coarse_codebook = np.load(coarse_path, mmap_mode="r")
# else:
# coarse_codebook = None
return ScannSearcher(
scann_pybind.ScannNumpy(db, train_set, scann_config, load_coarse,
coarse_path, load_fine, fine_path,
training_threads))
def load_searcher(artifacts_dir, assets_backcompat_shim=True):
"""Loads searcher assets from artifacts_dir and returns a ScaNN searcher."""
is_dir = os.path.isdir(artifacts_dir)
if not is_dir:
raise ValueError(f"{artifacts_dir} is not a directory.")
assets_pbtxt = os.path.join(artifacts_dir, "scann_assets.pbtxt")
if not scann_ops_pybind_backcompat.path_exists(assets_pbtxt):
if not assets_backcompat_shim:
raise ValueError("No scann_assets.pbtxt found.")
print("No scann_assets.pbtxt found. ScaNN assumes this searcher was from an"
" earlier release, and is calling `populate_and_save_assets_proto`"
"from `scann_ops_pybind_backcompat` to create a scann_assets.pbtxt. "
"Note this compatibility shim may be removed in the future.")
scann_ops_pybind_backcompat.populate_and_save_assets_proto(artifacts_dir)
with open(assets_pbtxt, "r") as f:
return ScannSearcher(scann_pybind.ScannNumpy(artifacts_dir, f.read()))
def create_searcher(db, scann_config, training_threads=0):
return ScannSearcher(
scann_pybind.ScannNumpy(db, scann_config, training_threads))
