INNER CODE UNIT · Python

pairwise_squared_euclidian_distance

openfoodfacts/openfoodfacts-ai · logo-ann/benchmarks/embedding_models_benchmark/main.py:120

def pairwise_squared_euclidian_distance(A: np.ndarray) -> torch.Tensor:
    assert len(A.shape) == 2
    dot_product = torch.matmul(A[:, None, :], A[None, :, :].swapaxes(1, 2)).squeeze()
    squared_sum = torch.sum(A**2.0, axis=1, keepdim=True)
    return squared_sum + squared_sum.T - 2 * dot_product


def pairwise_cosine_distance(A: torch.Tensor) -> torch.Tensor:
    assert len(A.shape) == 2
    normalized = torch.nn.functional.normalize(A, p=2.0, dim=1)
    return 1 - torch.matmul(normalized, normalized.T)


def run_model(
    root_dir: Path,
    split_set: Optional[Set[str]],
    model_name: str,
    batch_size: int,

View source record →