from enum import Enum
from typing import List, Optional, Union
import numpy as np
from qdrant_client.conversions import common_types as types
from qdrant_client.http import models
EPSILON = 1.1920929e-7 # https://doc.rust-lang.org/std/f32/constant.EPSILON.html
# https://github.com/qdrant/qdrant/blob/7164ac4a5987d28f1c93f5712aef8e09e7d93555/lib/segment/src/spaces/simple_avx.rs#L99C10-L99C10
class DistanceOrder(str, Enum):
BIGGER_IS_BETTER = "bigger_is_better"
SMALLER_IS_BETTER = "smaller_is_better"
class RecoQuery:
def __init__(
self,
positive: Optional[List[List[float]]] = None,
negative: Optional[List[List[float]]] = None,
):
positive = positive if positive is not None else []
negative = negative if negative is not None else []
self.positive: List[types.NumpyArray] = [np.array(vector) for vector in positive]
self.negative: List[types.NumpyArray] = [np.array(vector) for vector in negative]
assert not np.isnan(self.positive).any(), "Positive vectors must not contain NaN"
assert not np.isnan(self.negative).any(), "Negative vectors must not contain NaN"
class ContextPair:
def __init__(self, positive: List[float], negative: List[float]):
self.positive: types.NumpyArray = np.array(positive)
self.negative: types.NumpyArray = np.array(negative)
assert not np.isnan(self.positive).any(), "Positive vector must not contain NaN"
assert not np.isnan(self.negative).any(), "Negative vector must not contain NaN"
class DiscoveryQuery:
def __init__(self, target: List[float], context: List[ContextPair]):
self.target: types.NumpyArray = np.array(target)
self.context = context
assert not np.isnan(self.target).any(), "Target vector must not contain NaN"
class ContextQuery:
def __init__(self, context_pairs: List[ContextPair]):
self.context_pairs = context_pairs
DenseQueryVector = Union[
DiscoveryQuery,
ContextQuery,
RecoQuery,
]
def distance_to_order(distance: models.Distance) -> DistanceOrder:
"""
Convert distance to order
Args:
distance: distance to convert
Returns:
order
"""
if distance == models.Distance.EUCLID:
return DistanceOrder.SMALLER_IS_BETTER
elif distance == models.Distance.MANHATTAN:
return DistanceOrder.SMALLER_IS_BETTER
return DistanceOrder.BIGGER_IS_BETTER
def cosine_similarity(query: types.NumpyArray, vectors: types.NumpyArray) -> types.NumpyArray:
"""
Calculate cosine distance between query and vectors
Args:
query: query vector
vectors: vectors to calculate distance with
Returns:
distances
"""
vectors_norm = np.linalg.norm(vectors, axis=-1)[:, np.newaxis]
vectors /= np.where(vectors_norm != 0.0, vectors_norm, EPSILON)
if len(query.shape) == 1:
query_norm = np.linalg.norm(query)
query /= np.where(query_norm != 0.0, query_norm, EPSILON)
return np.dot(vectors, query)
query_norm = np.linalg.norm(query, axis=-1)[:, np.newaxis]
query /= np.where(query_norm != 0.0, query_norm, EPSILON)
return np.dot(query, vectors.T)
def dot_product(query: types.NumpyArray, vectors: types.NumpyArray) -> types.NumpyArray:
"""
Calculate dot product between query and vectors
Args:
query: query vector.
vectors: vectors to calculate distance with
Returns:
distances
"""
if len(query.shape) == 1:
return np.dot(vectors, query)
else:
return np.dot(query, vectors.T)
def euclidean_distance(query: types.NumpyArray, vectors: types.NumpyArray) -> types.NumpyArray:
"""
Calculate euclidean distance between query and vectors
Args:
query: query vector.
vectors: vectors to calculate distance with
Returns:
distances
"""
if len(query.shape) == 1:
return np.linalg.norm(vectors - query, axis=-1)
else:
return np.linalg.norm(vectors - query[:, np.newaxis], axis=-1)
def manhattan_distance(query: types.NumpyArray, vectors: types.NumpyArray) -> types.NumpyArray:
"""
Calculate manhattan distance between query and vectors
Args:
query: query vector.
vectors: vectors to calculate distance with
Returns:
distances
"""
if len(query.shape) == 1:
return np.sum(np.abs(vectors - query), axis=-1)
else:
return np.sum(np.abs(vectors - query[:, np.newaxis]), axis=-1)
def calculate_distance(
query: types.NumpyArray, vectors: types.NumpyArray, distance_type: models.Distance
) -> types.NumpyArray:
assert not np.isnan(query).any(), "Query vector must not contain NaN"
if distance_type == models.Distance.COSINE:
return cosine_similarity(query, vectors)
elif distance_type == models.Distance.DOT:
return dot_product(query, vectors)
elif distance_type == models.Distance.EUCLID:
return euclidean_distance(query, vectors)
elif distance_type == models.Distance.MANHATTAN:
return manhattan_distance(query, vectors)
else:
raise ValueError(f"Unknown distance type {distance_type}")
def calculate_distance_core(
query: types.NumpyArray, vectors: types.NumpyArray, distance_type: models.Distance
) -> types.NumpyArray:
"""
Calculate same internal distances as in core, rather than the final displayed distance
"""
assert not np.isnan(query).any(), "Query vector must not contain NaN"
if distance_type == models.Distance.EUCLID:
return -np.square(vectors - query, dtype=np.float32).sum(axis=1, dtype=np.float32)
if distance_type == models.Distance.MANHATTAN:
return -np.abs(vectors - query, dtype=np.float32).sum(axis=1, dtype=np.float32)
else:
return calculate_distance(query, vectors, distance_type)
def fast_sigmoid(x: np.float32) -> np.float32:
if np.isnan(x) or np.isinf(x):
# To avoid divisions on NaNs or inf, which gets: RuntimeWarning: invalid value encountered in scalar divide
return x
return x / np.add(1.0, abs(x))
def scaled_fast_sigmoid(x: np.float32) -> np.float32:
return 0.5 * (np.add(fast_sigmoid(x), 1.0))
def calculate_recommend_best_scores(
query: RecoQuery, vectors: types.NumpyArray, distance_type: models.Distance
) -> types.NumpyArray:
def get_best_scores(examples: List[types.NumpyArray]) -> types.NumpyArray:
vector_count = vectors.shape[0]
# Get scores to all examples
scores: List[types.NumpyArray] = []
for example in examples:
score = calculate_distance_core(example, vectors, distance_type)
scores.append(score)
# Keep only max for each vector
if len(scores) == 0:
scores.append(np.full(vector_count, -np.inf))
best_scores = np.array(scores, dtype=np.float32).max(axis=0)
return best_scores
pos = get_best_scores(query.positive)
neg = get_best_scores(query.negative)
# Choose from best positive or best negative,
# in in both cases we apply sigmoid and then negate depending on the order
return np.where(
pos > neg,
np.fromiter((scaled_fast_sigmoid(xi) for xi in pos), pos.dtype),
np.fromiter((-scaled_fast_sigmoid(xi) for xi in neg), neg.dtype),
)
def calculate_discovery_ranks(
context: List[ContextPair],
vectors: types.NumpyArray,
distance_type: models.Distance,
) -> types.NumpyArray:
overall_ranks = np.zeros(vectors.shape[0], dtype=np.int32)
for pair in context:
# Get distances to positive and negative vectors
pos = calculate_distance_core(pair.positive, vectors, distance_type)
neg = calculate_distance_core(pair.negative, vectors, distance_type)
pair_ranks = np.array(
[
1 if is_bigger else 0 if is_equal else -1
for is_bigger, is_equal in zip(pos > neg, pos == neg)
]
)
overall_ranks += pair_ranks
return overall_ranks
def calculate_discovery_scores(
query: DiscoveryQuery, vectors: types.NumpyArray, distance_type: models.Distance
) -> types.NumpyArray:
ranks = calculate_discovery_ranks(query.context, vectors, distance_type)
# Get distances to target
distances_to_target = calculate_distance_core(query.target, vectors, distance_type)
sigmoided_distances = np.fromiter(
(scaled_fast_sigmoid(xi) for xi in distances_to_target), np.float32
)
return ranks + sigmoided_distances
def calculate_context_scores(
query: ContextQuery, vectors: types.NumpyArray, distance_type: models.Distance
) -> types.NumpyArray:
overall_scores = np.zeros(vectors.shape[0], dtype=np.float32)
for pair in query.context_pairs:
# Get distances to positive and negative vectors
pos = calculate_distance_core(pair.positive, vectors, distance_type)
neg = calculate_distance_core(pair.negative, vectors, distance_type)
difference = pos - neg - EPSILON
pair_scores = np.fromiter(
(fast_sigmoid(xi) for xi in np.minimum(difference, 0.0)), np.float32
)
overall_scores += pair_scores
return overall_scores