Skip to content

Vector Normalization

Vector normalization utilities for GPU COSINE similarity support. Since GPU indexes don't support COSINE distance, we normalize vectors and use IP (Inner Product) distance instead.

NormalizingEmbeddings

Bases: Embeddings

Wrapper around an embedding model that automatically normalizes outputs. This is needed for GPU indexes when using COSINE similarity.

Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
class NormalizingEmbeddings(Embeddings):
    """
    Wrapper around an embedding model that automatically normalizes outputs.
    This is needed for GPU indexes when using COSINE similarity.
    """

    def __init__(self, embedding_model: Embeddings, normalize_for_gpu: bool = True):
        """
        Initialize the normalizing wrapper.

        Args:
            embedding_model: The underlying embedding model
            normalize_for_gpu: Whether to normalize embeddings (for GPU compatibility)
        """
        self.embedding_model = embedding_model
        self.normalize_for_gpu = normalize_for_gpu

        if normalize_for_gpu:
            logger.info(
                "Embedding model wrapped with normalization for GPU compatibility"
            )

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        """Embed documents and optionally normalize."""
        embeddings = self.embedding_model.embed_documents(texts)

        if self.normalize_for_gpu:
            embeddings = normalize_vectors_batch(embeddings)
            logger.debug("Normalized %d document embeddings for GPU", len(embeddings))

        return embeddings

    def embed_query(self, text: str) -> List[float]:
        """Embed query and optionally normalize."""
        embedding = self.embedding_model.embed_query(text)

        if self.normalize_for_gpu:
            embedding = normalize_vector(embedding)
            logger.debug("Normalized query embedding for GPU")

        return embedding

    def __getattr__(self, name):
        """Delegate other attributes to the underlying model."""
        return getattr(self.embedding_model, name)

__getattr__(name)

Delegate other attributes to the underlying model.

Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
112
113
114
def __getattr__(self, name):
    """Delegate other attributes to the underlying model."""
    return getattr(self.embedding_model, name)

__init__(embedding_model, normalize_for_gpu=True)

Initialize the normalizing wrapper.

Parameters:

Name Type Description Default
embedding_model Embeddings

The underlying embedding model

 required
normalize_for_gpu bool

Whether to normalize embeddings (for GPU compatibility)

 True
Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
def __init__(self, embedding_model: Embeddings, normalize_for_gpu: bool = True):
    """
    Initialize the normalizing wrapper.

    Args:
        embedding_model: The underlying embedding model
        normalize_for_gpu: Whether to normalize embeddings (for GPU compatibility)
    """
    self.embedding_model = embedding_model
    self.normalize_for_gpu = normalize_for_gpu

    if normalize_for_gpu:
        logger.info(
            "Embedding model wrapped with normalization for GPU compatibility"
        )

embed_documents(texts)

Embed documents and optionally normalize.

Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
 92
 93
 94
 95
 96
 97
 98
 99
100
def embed_documents(self, texts: List[str]) -> List[List[float]]:
    """Embed documents and optionally normalize."""
    embeddings = self.embedding_model.embed_documents(texts)

    if self.normalize_for_gpu:
        embeddings = normalize_vectors_batch(embeddings)
        logger.debug("Normalized %d document embeddings for GPU", len(embeddings))

    return embeddings

embed_query(text)

Embed query and optionally normalize.

Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
102
103
104
105
106
107
108
109
110
def embed_query(self, text: str) -> List[float]:
    """Embed query and optionally normalize."""
    embedding = self.embedding_model.embed_query(text)

    if self.normalize_for_gpu:
        embedding = normalize_vector(embedding)
        logger.debug("Normalized query embedding for GPU")

    return embedding

normalize_vector(vector)

Normalize a single vector to unit length.

Parameters:

Name Type Description Default
vector Union[List[float], ndarray]

Input vector as list or numpy array

 required

Returns:

Type Description
List[float]

Normalized vector as list
Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def normalize_vector(vector: Union[List[float], np.ndarray]) -> List[float]:
    """
    Normalize a single vector to unit length.

    Args:
        vector: Input vector as list or numpy array

    Returns:
        Normalized vector as list
    """
    vector = np.asarray(vector, dtype=np.float32)
    norm = np.linalg.norm(vector)

    if norm == 0:
        logger.warning("Zero vector encountered during normalization")
        return vector.tolist()

    normalized = vector / norm
    return normalized.tolist()

normalize_vectors_batch(vectors)

Normalize a batch of vectors to unit length.

Parameters:

Name Type Description Default
vectors List[List[float]]

List of vectors

 required

Returns:

Type Description
List[List[float]]

List of normalized vectors
Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
def normalize_vectors_batch(vectors: List[List[float]]) -> List[List[float]]:
    """
    Normalize a batch of vectors to unit length.

    Args:
        vectors: List of vectors

    Returns:
        List of normalized vectors
    """
    if not vectors:
        return vectors

    # Convert to numpy array for efficient computation
    vectors_array = np.asarray(vectors, dtype=np.float32)

    # Calculate norms for each vector
    norms = np.linalg.norm(vectors_array, axis=1, keepdims=True)

    # Handle zero vectors
    zero_mask = norms.flatten() == 0
    if np.any(zero_mask):
        logger.warning(
            "Found %d zero vectors during batch normalization", np.sum(zero_mask)
        )
        norms[zero_mask] = 1.0  # Avoid division by zero

    # Normalize
    normalized = vectors_array / norms

    return normalized.tolist()

should_normalize_vectors(has_gpu, use_cosine)

Determine if vectors should be normalized based on hardware and similarity metric.

Parameters:

Name Type Description Default
has_gpu bool

Whether GPU is being used

 required
use_cosine bool

Whether COSINE similarity is desired

 required

Returns:

Type Description
bool

True if vectors should be normalized
Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
def should_normalize_vectors(has_gpu: bool, use_cosine: bool) -> bool:
    """
    Determine if vectors should be normalized based on hardware and similarity metric.

    Args:
        has_gpu: Whether GPU is being used
        use_cosine: Whether COSINE similarity is desired

    Returns:
        True if vectors should be normalized
    """
    needs_normalization = has_gpu and use_cosine

    if needs_normalization:
        logger.info(
            "Vector normalization ENABLED: GPU detected with COSINE similarity request"
        )
    else:
        logger.info(
            "Vector normalization DISABLED: GPU=%s, COSINE=%s", has_gpu, use_cosine
        )

    return needs_normalization

wrap_embedding_model_if_needed(embedding_model, has_gpu, use_cosine=True)

Wrap embedding model with normalization if needed for GPU compatibility.

Parameters:

Name Type Description Default
embedding_model Embeddings

Original embedding model

 required
has_gpu bool

Whether GPU is being used

 required
use_cosine bool

Whether COSINE similarity is desired

 True

Returns:

Type Description
Embeddings

Original or wrapped embedding model
Source code in aiagents4pharma/talk2scholars/tools/pdf/utils/vector_normalization.py 
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def wrap_embedding_model_if_needed(
    embedding_model: Embeddings, has_gpu: bool, use_cosine: bool = True
) -> Embeddings:
    """
    Wrap embedding model with normalization if needed for GPU compatibility.

    Args:
        embedding_model: Original embedding model
        has_gpu: Whether GPU is being used
        use_cosine: Whether COSINE similarity is desired

    Returns:
        Original or wrapped embedding model
    """
    if should_normalize_vectors(has_gpu, use_cosine):
        return NormalizingEmbeddings(embedding_model, normalize_for_gpu=True)

    return embedding_model