fix: Fix get_embedding in the augmented classes

Author: xmnlab

poetry.lock

@@ -52,6 +52,7 @@ url = "https://download.pytorch.org/whl/cpu"
 priority = "explicit"
 
 [tool.poetry.dependencies]
+python = ">=3.9,<4"
 torch = [
   {version = ">=2.5.0", markers="extra!='gpu'", source="pytorch-cpu"},
   {version = ">=2.5.0", markers="extra=='gpu' and extra!='cpu'"},
@@ -88,7 +89,7 @@ makim = "1.19.0"
 virtualenv = "<=20.25.1"
 python-dotenv = ">=1.0"
 # note: Version 3.7.1 requries spaCy >=3.7.2,<3.8.0
-en-core-web-md = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_md-3.7.1/en_core_web_md-3.7.1-py3-none-any.whl"}
+en-core-web-md = {url = "https://github.com/explosion/spacy-models/releases/download/en_core_web_md-3.8.0/en_core_web_md-3.8.0-py3-none-any.whl"}
 
 [tool.pytest.ini_options]
 testpaths = [

pyproject.toml

@@ -18,6 +18,7 @@
 
 EmbeddingType: TypeAlias = Union[
     npt.NDArray[np.float64],
+    npt.NDArray[np.float32],
     Tensor,
     list[Tensor],
 ]

src/rago/augmented/base.py

@@ -38,8 +38,9 @@ def get_embedding(self, content: list[str]) -> EmbeddingType:
         response = model.embeddings.create(
             input=content, model=self.model_name
         )
-        result = np.array(response.data[0].embedding)
-        result = result.reshape(1, result.size)
+        result = np.array(
+            [data.embedding for data in response.data], dtype=np.float32
+        )
 
         self._save_cache(cache_key, result)
 

src/rago/augmented/openai.py

@@ -33,10 +33,22 @@ def get_embedding(self, content: List[str]) -> EmbeddingType:
 
         model = cast(spacy.language.Language, self.model)
         embeddings = []
+
         for text in content:
             doc = model(text)
+
+            # Ensure the model has proper vectors
+            if not doc.has_vector:
+                raise ValueError(f"Text: '{text}' has no valid word vectors!")
+
             embeddings.append(doc.vector)
-        result = np.array(embeddings)
+
+        result = np.array(embeddings, dtype=np.float32)
+
+        # Ensure 2D shape (num_texts, embedding_dim)
+        if result.ndim == 1:
+            result = result.reshape(1, -1)
+
         self._save_cache(cache_key, result)
         return result
 

src/rago/augmented/spacy.py

@@ -26,14 +26,26 @@ def api_key(env) -> str:
 
 
 @pytest.mark.skip_on_ci
-def test_aug_openai(animals_data: list[str], api_key: str) -> None:
+@pytest.mark.parametrize(
+    'question,expected_answer',
+    [
+        ('Is there any animal larger than a dinosaur?', 'Blue Whale'),
+        (
+            'What animal is renowned as the fastest animal on the planet?',
+            'Peregrine Falcon',
+        ),
+        ('An animal which do pollination?', 'Honey Bee'),
+    ],
+)
+def test_aug_openai(
+    animals_data: list[str], api_key: str, question: str, expected_answer: str
+) -> None:
     """Test RAG pipeline with OpenAI's GPT."""
     logs = {
         'augmented': {},
     }
 
-    query = 'Is there any animal larger than a dinosaur?'
-    top_k = 3
+    top_k = 2
 
     ret_string = StringRet(animals_data)
     aug_openai = OpenAIAug(
@@ -43,14 +55,14 @@ def test_aug_openai(animals_data: list[str], api_key: str) -> None:
     )
 
     ret_result = ret_string.get()
-    aug_result = aug_openai.search(query, ret_result)
+    aug_result = aug_openai.search(question, ret_result)
 
     assert aug_openai.top_k == top_k
     # note: openai as augmented doesn't work as expected
     #   it is returning a very poor result
     #   it needs to be revisited and improved
     assert len(aug_result) >= 1
-    assert 'blue whale' in aug_result[0].lower()
+    assert expected_answer.lower() in aug_result[0].lower()
 
     # check if logs have been used
     assert logs['augmented']

tests/test_openai.py

@@ -1,30 +1,44 @@
 """Tests for Rago package using SpaCy."""
 
+import pytest
+
 from rago.augmented import SpaCyAug
-from rago.retrieval import StringRet
 
 
-def test_aug_spacy(animals_data: list[str]) -> None:
+@pytest.mark.parametrize(
+    'question,expected_answer',
+    [
+        ('Is there any animal larger than a dinosaur?', 'Blue Whale'),
+        (
+            'What animal is renowned as the fastest animal on the planet?',
+            'Peregrine Falcon',
+        ),
+        ('An animal which do pollination?', 'Honey Bee'),
+    ],
+)
+def test_aug_spacy(
+    animals_data: list[str], question: str, expected_answer: str
+) -> None:
     """Test RAG pipeline with SpaCy."""
     logs = {
         'augmented': {},
     }
 
-    query = 'Is there any animal larger than a dinosaur?'
-    top_k = 3
+    top_k = 2
 
-    ret_string = StringRet(animals_data)
     aug_openai = SpaCyAug(
+        model_name='en_core_web_md',
         top_k=top_k,
         logs=logs['augmented'],
     )
 
-    ret_result = ret_string.get()
-    aug_result = aug_openai.search(query, ret_result)
+    aug_result = aug_openai.search(question, animals_data)
 
     assert aug_openai.top_k == top_k
-    assert len(aug_result) == top_k
-    assert any(['blue whale' in result.lower() for result in aug_result])
+    assert top_k >= len(aug_result)
+    assert any(
+        [expected_answer.lower() in result.lower() for result in aug_result]
+    )
 
     # check if logs have been used
     assert logs['augmented']