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))