Added handling edge cases in scorer.py

app/scorer.py

@@ -1,9 +1,25 @@
 import numpy as np
 
 
+# TODO: Move exception classes into separate file
+class DimensionalityMismatchError(ValueError):
+    """Raised when the dimensions of query and corpus vectors don't match."""
+    pass
+
+
+class ZeroVectorError(ValueError):
+    """Raised when a zero vector is encountered."""
+    pass
+
+
+class EmptyInputError(ValueError):
+    """Raised when the input arrays are empty."""
+    pass
+
+
 def cosine_similarity(
-     query_vector: np.ndarray,
-     corpus_vectors: np.ndarray
+        query_vector: np.ndarray,
+        corpus_vectors: np.ndarray
 ) -> np.ndarray:
     """
     Calculate cosine similarity between prompt vectors.
@@ -12,9 +28,39 @@ def cosine_similarity(
         corpus_vectors: Vectorized prompt corpus of shape (N, D).
     Returns:
         The vector of shape (N,) with values in range [-1, 1] where 1 is max similarity i.e., two vectors are the same.
+    Raises:
+        DimensionalityMismatchError: If dimensions of query_vector and corpus_vectors do not match.
+        ZeroVectorError: If query_vector is a zero vector or any corpus vector is a zero vector.
+        EmptyInputError: If query_vector or corpus_vectors are empty.
+
+    Note:
+        - This implementation assumes the use of SentenceTransformer with the "all-MiniLM-L6-v2" model.
+        - SentenceTransformer embeddings are unlikely to produce zero vectors, even for empty or irrelevant inputs.
+        - However, checks for zero vectors are included to handle potential edge cases and ensure robustness
+          for future modifications or alternative embedding models.
     """
+    # Validate input shapes and properties
+    if query_vector.shape[0] != 1:
+        raise DimensionalityMismatchError(f"query_vector must have shape (1, D), but got shape {query_vector.shape}.")
+    if query_vector.shape[1] != corpus_vectors.shape[1]:
+        raise DimensionalityMismatchError(
+            f"query_vector shape {query_vector.shape} does not match corpus_vectors shape {corpus_vectors.shape}."
+        )
+    if query_vector.size == 0 or corpus_vectors.size == 0:
+        raise EmptyInputError("query_vector and corpus_vectors must not be empty.")
+
+    # Compute query norm and check for zero vector
     query_norm = np.linalg.norm(query_vector, axis=1)[0]
+    if query_norm == 0:
+        raise ZeroVectorError("query_vector must not be a zero vector.")
+
+    # Check if any corpus vector is zero (alternative option for raising an error if any corpus vector is zero,
+    # can be filtering out zero vectors and raising the error only if all corpus vectors are zero vectors)
     corpus_norms = np.linalg.norm(corpus_vectors, axis=1)
+    if np.any(corpus_norms == 0):
+        raise ZeroVectorError("corpus_vectors must not contain zero vectors.")
+
+    # Compute cosine similarity
     dot_products = np.dot(corpus_vectors, query_vector.T).flatten()
     similarities = dot_products / (query_norm * corpus_norms)
     return similarities