fix: evaluation pipeline should return all element types (#242)

deepsense-ai · Dec 9, 2024 · 32ba29c · 32ba29c
1 parent 439e445
commit 32ba29c
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/packages/ragbits-evaluate/CHANGELOG.md b/packages/ragbits-evaluate/CHANGELOG.md
@@ -2,6 +2,10 @@
 
 ## Unreleased
 
+### Changed
+
+- document search evaluation now returns all Element types, rather than only TextElements (#241).
+
 ## 0.5.0 (2024-12-05)
 
 ### Changed

diff --git a/packages/ragbits-evaluate/src/ragbits/evaluate/pipelines/document_search.py b/packages/ragbits-evaluate/src/ragbits/evaluate/pipelines/document_search.py
@@ -8,7 +8,6 @@
 
 from ragbits.document_search import DocumentSearch
 from ragbits.document_search.documents.document import DocumentMeta
-from ragbits.document_search.documents.element import TextElement
 from ragbits.document_search.documents.sources import HuggingFaceSource
 from ragbits.evaluate.pipelines.base import EvaluationPipeline, EvaluationResult
 
@@ -50,7 +49,7 @@ async def __call__(self, data: dict) -> DocumentSearchResult:
             The evaluation result.
         """
         elements = await self.document_search.search(data["question"])
-        predicted_passages = [element.content for element in elements if isinstance(element, TextElement)]
+        predicted_passages = [element.text_representation or "" for element in elements]
         return DocumentSearchResult(
             question=data["question"],
             reference_passages=data["passages"],