feat: сhanged error handling in the statement extractor (#61)

Author: Efreet408

src/aidial_rag_eval/generation/inference.py

@@ -224,7 +224,7 @@ def segment_hypotheses(
 
 
 def extract_statements(
-    hypotheses_segments: List[List[HypothesisSegment]],
+    list_of_hypothesis_segments: List[List[HypothesisSegment]],
     llm: BaseChatModel,
     max_concurrency: int = 8,
     show_progress_bar: bool = True,
@@ -237,7 +237,7 @@ def extract_statements(
     Parameters
     -----------
 
-    hypotheses_segments : List[List[HypothesisSegment]]
+    list_of_hypothesis_segments : List[List[HypothesisSegment]]
         Nested list of hypothesis segments.
 
     llm : BaseChatModel
@@ -264,7 +264,7 @@ def extract_statements(
     if show_progress_bar:
         print("Extracting statements...")
     statements = extractor.extract(
-        hypotheses_segments,
+        list_of_hypothesis_segments,
         show_progress_bar,
     )
     return statements
@@ -404,7 +404,7 @@ def calculate_batch_inference(
         show_progress_bar=show_progress_bar,
     )
     statements: List[List[List[Statement]]] = extract_statements(
-        hypotheses_segments=[
+        list_of_hypothesis_segments=[
             segmented_hypothesis.segments
             for segmented_hypothesis in segmented_hypotheses
         ],

src/aidial_rag_eval/generation/models/statement_extractor/base_statement_extractor.py

@@ -14,7 +14,7 @@ class StatementExtractor(ABC):
     @abstractmethod
     def extract(
         self,
-        hypothesis_segments: List[List[HypothesisSegment]],
+        list_of_hypothesis_segments: List[List[HypothesisSegment]],
         show_progress_bar: bool,
     ) -> List[List[List[Statement]]]:
         pass

src/aidial_rag_eval/generation/models/statement_extractor/llm_statement_extractor.py

@@ -1,7 +1,7 @@
 from typing import Dict, List
 
 from langchain_core.language_models import BaseChatModel
-from langchain_core.runnables import RunnableSerializable, chain
+from langchain_core.runnables import RunnablePassthrough, RunnableSerializable, chain
 
 from aidial_rag_eval.generation.models.lambdas import json_to_list
 from aidial_rag_eval.generation.models.statement_extractor.base_statement_extractor import (
@@ -16,40 +16,51 @@
 
 @chain
 def list_to_statements(
-    statements_for_each_hypothesis: List[Dict[str, List[str]]],
+    llm_outputs_with_inputs: Dict,
 ) -> List[List[str]]:
     """
     Function is part of a chain that extracts segments from a list.
 
     Parameters
     -----------
-    statements_for_each_hypothesis : List[Dict[str, List[str]]]
-        The output list of dicts from the LLM with extracted statements.
+    llm_outputs_with_inputs : Dict
+        Passed inputs with output list of dicts from the LLM with extracted statements.
 
     Returns
     ------------
-    List[str]
+    List[List[str]]
         The extracted statements. if the LLM output is valid;
-        otherwise, an empty list is returned.
+        otherwise, return an original list of hypothesis_segments,
+        each hypothesis_segment is wrapped.
     """
     try:
+        statements_for_hypothesis_segments = llm_outputs_with_inputs[
+            "llm_output_statements"
+        ]
+        hypothesis_segments = llm_outputs_with_inputs["hypothesis_segments"]
+        assert len(hypothesis_segments) == len(statements_for_hypothesis_segments)
         return [
-            return_dict["statements"] for return_dict in statements_for_each_hypothesis
+            return_dict["statements"]
+            for return_dict in llm_outputs_with_inputs["llm_output_statements"]
         ]
     except (
         TypeError,
         KeyError,
+        AssertionError,
     ):
-        return []
+        return [
+            [hypothesis_segment]
+            for hypothesis_segment in llm_outputs_with_inputs["hypothesis_segments"]
+        ]
 
 
 @chain
 def wrap_hypotheses(input_: Dict) -> Dict:
     assert type(input_) is dict
     return {
         "hypotheses": [
-            f"<hypothesis{index + 1}> {hypothesis} </hypothesis{index + 1}>"
-            for index, hypothesis in enumerate(input_["hypotheses"])
+            f"<hypothesis{index + 1}> {hypothesis_segment} </hypothesis{index + 1}>"
+            for index, hypothesis_segment in enumerate(input_["hypothesis_segments"])
         ],
     }
 
@@ -76,17 +87,19 @@ def __init__(
     ):
 
         self._chain = (
-            wrap_hypotheses
-            | statement_prompt
-            | model
-            | json_to_list
+            RunnablePassthrough.assign(
+                llm_output_statements=wrap_hypotheses
+                | statement_prompt
+                | model
+                | json_to_list
+            )
             | list_to_statements
         )
         self.max_concurrency = max_concurrency
 
     def extract(
         self,
-        hypothesis_segments: List[List[HypothesisSegment]],
+        list_of_hypothesis_segments: List[List[HypothesisSegment]],
         show_progress_bar: bool,
     ) -> List[List[List[Statement]]]:
         """
@@ -95,7 +108,7 @@ def extract(
 
         Parameters
         -----------
-        hypothesis_segments : List[List[HypothesisSegment]]
+        list_of_hypothesis_segments : List[List[HypothesisSegment]]
             A list of hypothesis segments as a sources of statements.
 
         show_progress_bar : bool
@@ -107,13 +120,15 @@ def extract(
             Returns the statements for each hypothesis segment.
         """
 
-        with ProgressBarCallback(len(hypothesis_segments), show_progress_bar) as cb:
+        with ProgressBarCallback(
+            len(list_of_hypothesis_segments), show_progress_bar
+        ) as cb:
             returns = self._chain.batch(
                 [
                     {
-                        "hypotheses": hypotheses,
+                        "hypothesis_segments": hypothesis_segments,
                     }
-                    for hypotheses in hypothesis_segments
+                    for hypothesis_segments in list_of_hypothesis_segments
                 ],
                 config={"callbacks": [cb], "max_concurrency": self.max_concurrency},
             )