feat: decontextualization returns new objects without mutating inputs (#63)

Author: Efreet408

src/aidial_rag_eval/generation/inference.py

@@ -219,8 +219,10 @@ def segment_hypotheses(
     ]
     if show_progress_bar:
         print("Converting hypothesis...")
-    converter.transform_texts(segmented_hypotheses, show_progress_bar)
-    return segmented_hypotheses
+    decontextualized__segmented_hypotheses = converter.transform_texts(
+        segmented_hypotheses, show_progress_bar
+    )
+    return decontextualized__segmented_hypotheses
 
 
 def extract_statements(

src/aidial_rag_eval/generation/models/converters/base_converter.py

@@ -14,5 +14,5 @@ class SegmentConverter(ABC):
     @abstractmethod
     def transform_texts(
         self, segmented_texts: List[SegmentedText], show_progress_bar: bool
-    ):
+    ) -> List[SegmentedText]:
         pass

src/aidial_rag_eval/generation/models/converters/llm_decontextualization_converter.py

@@ -5,18 +5,28 @@
 from langchain_core.exceptions import OutputParserException
 from langchain_core.language_models import BaseChatModel
 from langchain_core.messages import AIMessage
-from langchain_core.runnables import RunnableSerializable, chain
+from langchain_core.runnables import (
+    RunnableBranch,
+    RunnablePassthrough,
+    RunnableSerializable,
+    chain,
+)
 from langchain_core.utils.json import parse_json_markdown
 
 from aidial_rag_eval.generation.models.converters.base_converter import SegmentConverter
 from aidial_rag_eval.generation.models.converters.decontextualization_template import (
     decontextualization_prompt,
 )
-from aidial_rag_eval.generation.types import TextSegment
 from aidial_rag_eval.generation.utils.progress_bar import ProgressBarCallback
 from aidial_rag_eval.generation.utils.segmented_text import SegmentedText
 
 
+@chain
+def check_if_sentences_less_than_2(input_: Dict) -> bool:
+    assert type(input_) is dict
+    return len(input_["segmented_text"].segments) < 2
+
+
 @chain
 def json_to_dict_segments(input_: AIMessage) -> List[str]:
     """
@@ -48,21 +58,40 @@ def json_to_dict_segments(input_: AIMessage) -> List[str]:
 
 
 @chain
-def sentences_to_json_list(input_: Dict) -> Dict:
+def segmented_text_to_json_list(input_: Dict) -> Dict:
+    assert type(input_) is dict
+    return {"sentences_str": json.dumps(input_["segmented_text"].segments)}
+
+
+@chain
+def return_original_segmented_text(input_: Dict) -> Dict:
     assert type(input_) is dict
-    return {"sentences_str": json.dumps(input_["sentences"])}
+    return input_["segmented_text"]
+
+
+@chain
+def dict_segments_to_segmented_text(llm_outputs_with_inputs: Dict) -> SegmentedText:
+    original_segmented_text: SegmentedText = llm_outputs_with_inputs["segmented_text"]
+    try:
+        decontextualized_segments = llm_outputs_with_inputs["decontextualized_segments"]
+        assert len(decontextualized_segments) == len(original_segmented_text.segments)
+        return SegmentedText(
+            decontextualized_segments, original_segmented_text.delimiters
+        )
+    except (TypeError, KeyError, AssertionError):
+        return original_segmented_text
 
 
 class LLMNoPronounsConverter(SegmentConverter):
     """
-    The LLMNoPronounsBatchConverter is designed to replace pronouns
-    in text segments using a LLM.
-
-    Input is a list of SegmentedText objects.
-    If a SegmentedText object contains more than one segment,
-    segments are sent in a prompt to the LLM.
-    In a prompt, the first segment is used only for context,
-    and pronoun replacement is performed only in the remaining segments.
+    Converter that decontextualizes text segments using an LLM.
+
+    Takes a list of SegmentedText objects and processes each one:
+    - If a SegmentedText has fewer than 2 segments, it is returned unchanged.
+    - Otherwise, all segments are sent to the LLM for decontextualization.
+
+    The LLM replaces pronouns and context-dependent references
+    to make each segment self-contained.
     """
 
     _chain: RunnableSerializable
@@ -79,24 +108,24 @@ def __init__(
         model: BaseChatModel,
         max_concurrency: int,
     ):
-
-        self._chain = (
-            sentences_to_json_list
-            | decontextualization_prompt
-            | model
-            | json_to_dict_segments
+        self._chain = RunnableBranch(
+            (check_if_sentences_less_than_2, return_original_segmented_text),
+            RunnablePassthrough.assign(
+                decontextualized_segments=segmented_text_to_json_list
+                | decontextualization_prompt
+                | model
+                | json_to_dict_segments
+            )
+            | dict_segments_to_segmented_text,
         )
         self.max_concurrency = max_concurrency
 
     def transform_texts(
         self, segmented_texts: List[SegmentedText], show_progress_bar: bool
-    ):
+    ) -> List[SegmentedText]:
         """
         Method that converts segmented texts by replacing pronouns using an LLM.
-        The LLM processes segments,
-        where the additional first segment is not converted
-        but is provided for context to enable the conversion of the second sentence.
-        The LLM returns converted segments.
+        The LLM processes all segments and returns converted segments.
         If the invariant of the length of input and output segment batches
         is not maintained, the segments of this batch are not replaced.
 
@@ -107,60 +136,20 @@ def transform_texts(
 
         show_progress_bar : bool
             A flag that controls the display of a progress bar.
-        """
-        original_segment_batches: List[List[TextSegment]] = []
-        segment_ids: List[int] = []
-        for text_id, segmented_text in enumerate(segmented_texts):
-            segments = segmented_text.segments
-            if len(segments) <= 1:
-                continue
-            original_segment_batches.append(segments)
-            segment_ids.append(text_id)
-
-        no_pronouns_segment_batches = self._get_no_pronouns_segments(
-            original_segment_batches, show_progress_bar
-        )
-
-        for text_id, no_pronouns_segment_batch, original_segment_batch in zip(
-            segment_ids, no_pronouns_segment_batches, original_segment_batches
-        ):
-            if len(no_pronouns_segment_batch) != len(original_segment_batch):
-                continue
-            segmented_texts[text_id].replace_segments(
-                no_pronouns_segment_batch[1:],
-                1,
-            )
 
-    def _get_no_pronouns_segments(
-        self,
-        original_segment_batches: List[List[TextSegment]],
-        show_progress_bar: bool,
-    ) -> List[List[TextSegment]]:
-        """
-        Method that calls _chain to replace pronouns.
-
-        Parameters
-        -----------
-        original_segment_batches : List[List[str]]
-            Segments of texts.
-
-        show_progress_bar : bool
-            A flag that controls the display of a progress bar.
         Returns
-        ------------
-        List[List[str]]
-            List of converted segments, divided into batches.
+        -------
+        List[SegmentedText]
+            A list of segmented texts with decontextualized segments.
         """
-        with ProgressBarCallback(
-            len(original_segment_batches), show_progress_bar
-        ) as cb:
-            no_pronouns_segment_batches = self._chain.batch(
+        with ProgressBarCallback(len(segmented_texts), show_progress_bar) as cb:
+            decontextualized_segmented_texts = self._chain.batch(
                 [
                     {
-                        "sentences": batch,
+                        "segmented_text": segmented_text,
                     }
-                    for batch in original_segment_batches
+                    for segmented_text in segmented_texts
                 ],
                 config={"callbacks": [cb], "max_concurrency": self.max_concurrency},
             )
-        return no_pronouns_segment_batches
+        return decontextualized_segmented_texts

tests/chain_tests/test_decontextualization_chain.py

@@ -18,9 +18,14 @@ def test_valid_json_response():
         segments=["John went to the store.", "He bought milk."], delimiters=[" "]
     )
 
-    converter.transform_texts([segmented_text], show_progress_bar=False)
+    decontext_segmented_text = converter.transform_texts(
+        [segmented_text], show_progress_bar=False
+    )[0]
 
-    assert segmented_text.segments == ["John went to the store.", "John bought milk."]
+    assert decontext_segmented_text.segments == [
+        "John went to the store.",
+        "John bought milk.",
+    ]
 
 
 def test_invalid_json_response():
@@ -30,11 +35,12 @@ def test_invalid_json_response():
     segmented_text = SegmentedText(
         segments=["John went to the store.", "He bought milk."], delimiters=[" "]
     )
-    original_segments = segmented_text.segments.copy()
 
-    converter.transform_texts([segmented_text], show_progress_bar=False)
+    decontext_segmented_text = converter.transform_texts(
+        [segmented_text], show_progress_bar=False
+    )[0]
 
-    assert segmented_text.segments == original_segments
+    assert decontext_segmented_text.segments == segmented_text.segments
 
 
 def test_json_missing_segments_key():
@@ -46,11 +52,12 @@ def test_json_missing_segments_key():
     segmented_text = SegmentedText(
         segments=["John went to the store.", "He bought milk."], delimiters=[" "]
     )
-    original_segments = segmented_text.segments.copy()
 
-    converter.transform_texts([segmented_text], show_progress_bar=False)
+    decontext_segmented_text = converter.transform_texts(
+        [segmented_text], show_progress_bar=False
+    )[0]
 
-    assert segmented_text.segments == original_segments
+    assert decontext_segmented_text.segments == segmented_text.segments
 
 
 def test_segment_count_mismatch():
@@ -60,11 +67,12 @@ def test_segment_count_mismatch():
     segmented_text = SegmentedText(
         segments=["John went to the store.", "He bought milk."], delimiters=[" "]
     )
-    original_segments = segmented_text.segments.copy()
 
-    converter.transform_texts([segmented_text], show_progress_bar=False)
+    decontext_segmented_text = converter.transform_texts(
+        [segmented_text], show_progress_bar=False
+    )[0]
 
-    assert segmented_text.segments == original_segments
+    assert decontext_segmented_text.segments == segmented_text.segments
 
 
 def test_empty_response():
@@ -74,11 +82,12 @@ def test_empty_response():
     segmented_text = SegmentedText(
         segments=["John went to the store.", "He bought milk."], delimiters=[" "]
     )
-    original_segments = segmented_text.segments.copy()
 
-    converter.transform_texts([segmented_text], show_progress_bar=False)
+    decontext_segmented_text = converter.transform_texts(
+        [segmented_text], show_progress_bar=False
+    )[0]
 
-    assert segmented_text.segments == original_segments
+    assert decontext_segmented_text.segments == segmented_text.segments
 
 
 def test_invoke_raises_exception():