Function load_from_hub

vicinity/integrations/huggingface.py:101–132

Load a Vicinity instance from the Hugging Face Hub. :param repo_id: The repository ID on the Hugging Face Hub. :param token: Optional authentication token for private repositories. :param **kwargs: Additional arguments passed to load_dataset. :return: A Vicinity instance loaded

(
    repo_id: str, token: str | None = None, **kwargs: Any
)

def load_from_hub(
    repo_id: str, token: str | None = None, **kwargs: Any
) -> tuple[list[dict[str, Any]], BasicVectorStore | None, AbstractBackend, dict[str, Any]]:
    """
    Load a Vicinity instance from the Hugging Face Hub.

    :param repo_id: The repository ID on the Hugging Face Hub.
    :param token: Optional authentication token for private repositories.
    :param **kwargs: Additional arguments passed to load_dataset.
    :return: A Vicinity instance loaded from the Hub.
    """
    # Load dataset and extract items and vectors
    dataset = load_dataset(repo_id, token=token, split="train", **kwargs)
    has_vectors = "vectors" in dataset.column_names
    if has_vectors:
        vectors = dataset["vectors"]
    if _VICINITY_ITEM_COLUMN in dataset.column_names:
        items = dataset[_VICINITY_ITEM_COLUMN]
    else:
        # Remove the vectors column (if it exists) and then convert to a list.
        if has_vectors:
            dataset = dataset.remove_columns("vectors")
        items = dataset.to_list()
    vector_store = BasicVectorStore(vectors=vectors) if has_vectors else None

    # Download and load config and backend
    repo_path = Path(snapshot_download(repo_id=repo_id, token=token, repo_type="dataset"))
    with open(repo_path / "config.json") as f:
        config = json.load(f)
    backend_type = Backend(config["backend_type"])
    backend = get_backend_class(backend_type).load(repo_path / "backend")
    return items, vector_store, backend, config