From 93542e8f7a1f340b1609f59c31e0c3105ab07367 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 25 Feb 2025 19:34:12 -0800
Subject: [PATCH 01/12] add cot support for sem_map

---
 lotus/sem_ops/postprocessors.py      |  43 ++--------
 lotus/sem_ops/sem_map.py             |  27 +++++-
 lotus/templates/task_instructions.py | 120 +++++++++++----------------
 3 files changed, 79 insertions(+), 111 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index a361a166..ea6eb4b7 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -28,51 +28,24 @@ def cot_postprocessor(llm_answers: list[str]):
     return outputs, explanations
 
 
-def map_postprocess_cot(llm_answers: list[str]) -> SemanticMapPostprocessOutput:
-    """
-    Postprocess the output of the map operator with CoT reasoning.
-
-    Args:
-        llm_answers (list[str]): The list of llm answers.
-
-    Returns:
-        SemanticMapPostprocessOutput
-    """
-    outputs: list[str] = []
-    explanations: list[str | None] = []
-
-    for llm_answer in llm_answers:
-        reasoning_idx = llm_answer.find("Reasoning:\n")
-        if reasoning_idx == -1:
-            reasoning_idx = 0
-        else:
-            reasoning_idx += len("Reasoning:\n")
-
-        answer_idx = llm_answer.find("Answer:")
-        reasoning = llm_answer[reasoning_idx:answer_idx].rstrip("\n").lstrip("\n")
-        answer = llm_answer[answer_idx + len("Answer:") :]
-        outputs.append(answer)
-        explanations.append(reasoning)
-
-    return SemanticMapPostprocessOutput(raw_outputs=llm_answers, outputs=outputs, explanations=explanations)
-
-
-def map_postprocess(llm_answers: list[str], cot_reasoning: bool = False) -> SemanticMapPostprocessOutput:
+def map_postprocess(llm_answers: list[str], default: str = "") -> SemanticMapPostprocessOutput:
     """
     Postprocess the output of the map operator.
 
     Args:
         llm_answers (list[str]): The list of llm answers.
-        cot_reasoning (bool): Whether there is CoT reasoning.
+        default (str): The default value to use if we fail to parse the answer.
 
     Returns:
         SemanticMapPostprocessOutput
     """
-    if cot_reasoning:
-        return map_postprocess_cot(llm_answers)
+    outputs, explanations = cot_postprocessor(llm_answers)
+
+    for i, output in enumerate(outputs):
+        if output is None:
+            lotus.logger.info(f"\t Failed to parse {output}: defaulting to {default}")
+            outputs[i] = default
 
-    outputs: list[str] = llm_answers
-    explanations: list[str | None] = [None] * len(llm_answers)
     return SemanticMapPostprocessOutput(raw_outputs=llm_answers, outputs=outputs, explanations=explanations)
 
 
diff --git a/lotus/sem_ops/sem_map.py b/lotus/sem_ops/sem_map.py
index 4cc22d88..508f4a82 100644
--- a/lotus/sem_ops/sem_map.py
+++ b/lotus/sem_ops/sem_map.py
@@ -15,13 +15,15 @@ def sem_map(
     docs: list[dict[str, Any]],
     model: lotus.models.LM,
     user_instruction: str,
-    postprocessor: Callable[[list[str], bool], SemanticMapPostprocessOutput] = map_postprocess,
+    default: str = "",
+    postprocessor: Callable[[list[str]], SemanticMapPostprocessOutput] | None = None,
     examples_multimodal_data: list[dict[str, Any]] | None = None,
     examples_answers: list[str] | None = None,
     cot_reasoning: list[str] | None = None,
     strategy: str | None = None,
     safe_mode: bool = False,
     progress_bar_desc: str = "Mapping",
+    additional_cot_instructions: str = "",
 ) -> SemanticMapOutput:
     """
     Maps a list of documents to a list of outputs using a model.
@@ -42,7 +44,13 @@ def sem_map(
     inputs = []
     for doc in docs:
         prompt = lotus.templates.task_instructions.map_formatter(
-            doc, user_instruction, examples_multimodal_data, examples_answers, cot_reasoning, strategy=strategy
+            doc,
+            user_instruction,
+            examples_multimodal_data,
+            examples_answers,
+            cot_reasoning,
+            strategy=strategy,
+            reasoning_instructions=additional_cot_instructions,
         )
         lotus.logger.debug(f"input to model: {prompt}")
         lotus.logger.debug(f"inputs content to model: {[x.get('content') for x in prompt]}")
@@ -58,7 +66,11 @@ def sem_map(
     lm_output: LMOutput = model(inputs, progress_bar_desc=progress_bar_desc)
 
     # post process results
-    postprocess_output = postprocessor(lm_output.outputs, strategy in ["cot", "zs-cot"])
+    if postprocessor:
+        postprocess_output = postprocessor(lm_output.outputs)
+    else:
+        postprocess_output = map_postprocess(lm_output.outputs, default=default)
+
     lotus.logger.debug(f"raw_outputs: {lm_output.outputs}")
     lotus.logger.debug(f"outputs: {postprocess_output.outputs}")
     lotus.logger.debug(f"explanations: {postprocess_output.explanations}")
@@ -89,7 +101,7 @@ def _validate(obj: pd.DataFrame) -> None:
     def __call__(
         self,
         user_instruction: str,
-        postprocessor: Callable[[list[str], bool], SemanticMapPostprocessOutput] = map_postprocess,
+        postprocessor: Callable[[list[str]], SemanticMapPostprocessOutput] | None = None,
         return_explanations: bool = False,
         return_raw_outputs: bool = False,
         suffix: str = "_map",
@@ -97,6 +109,8 @@ def __call__(
         strategy: str | None = None,
         safe_mode: bool = False,
         progress_bar_desc: str = "Mapping",
+        additional_cot_instructions: str = "",
+        default: str = "",
     ) -> pd.DataFrame:
         """
         Applies semantic map over a dataframe.
@@ -109,6 +123,9 @@ def __call__(
             suffix (str): The suffix for the new columns. Defaults to "_map".
             examples (pd.DataFrame | None): The examples dataframe. Defaults to None.
             strategy (str | None): The reasoning strategy. Defaults to None.
+            safe_mode (bool): Whether to use safe mode. Defaults to False.
+            progress_bar_desc (str): The description for the progress bar. Defaults to "Mapping".
+            additional_cot_instructions (str): Additional instructions for the CoT. Defaults to "".
 
         Returns:
             pd.DataFrame: The dataframe with the new mapped columns.
@@ -151,6 +168,8 @@ def __call__(
             strategy=strategy,
             safe_mode=safe_mode,
             progress_bar_desc=progress_bar_desc,
+            additional_cot_instructions=additional_cot_instructions,
+            default=default,
         )
 
         new_df = self._obj.copy()
diff --git a/lotus/templates/task_instructions.py b/lotus/templates/task_instructions.py
index a71acd8c..ebe6387e 100644
--- a/lotus/templates/task_instructions.py
+++ b/lotus/templates/task_instructions.py
@@ -9,25 +9,25 @@
 
 
 def cot_formatter(reasoning, answer):
-    return f"""Reasoning:\n{reasoning}\n\nAnswer: {answer}"""
+    return f"""\n\nReasoning:\n{reasoning}\n\nAnswer: {answer}"""
 
 
 def answer_only_formatter(answer):
-    return f"""Answer: {answer}"""
+    return f"""\n\nAnswer: {answer}"""
 
 
 def cot_prompt_formatter(reasoning_instructions: str = "", answer_instructions: str = "") -> str:
-    reasoning_instructions = f"<Your reasoning here. {reasoning_instructions}>"
-    answer_instructions = f"<Your answer here. {answer_instructions}>"
-    return f"""Let's think step by step. Use the following format to provide your answer:
-        {cot_formatter(reasoning_instructions, answer_instructions)}
+    reasoning_placeholder = f"<Your reasoning here. {reasoning_instructions}>"
+    answer_placeholder = f"<Your answer here. {answer_instructions}>"
+    return f"""\n\nLet's think step by step. Use the following format to provide your answer:
+        {cot_formatter(reasoning_placeholder, answer_placeholder)}
         """
 
 
 def non_cot_prompt_formatter(answer_instructions: str = "") -> str:
-    answer_instructions = f"<Your answer here. {answer_instructions}>"
-    return f"""Use the following format to provide your answer:
-            {answer_only_formatter(answer_instructions)}
+    answer_placeholder = f"<Your answer here. {answer_instructions}>"
+    return f"""\n\nUse the following format to provide your answer:
+            {answer_only_formatter(answer_placeholder)}
             """
 
 
@@ -143,57 +143,6 @@ def filter_formatter(
     return messages
 
 
-def map_formatter_cot(
-    multimodal_data: dict[str, Any],
-    user_instruction: str,
-    examples_multimodal_data: list[dict[str, Any]],
-    examples_answer: list[str],
-    cot_reasoning: list[str],
-) -> list[dict[str, str]]:
-    sys_instruction = (
-        "The user will provide an instruction and some relevant context.\n"
-        "Your job is to answer the user's instruction given the context."
-        "You must give your reasoning and then your final answer"
-    )
-    messages = [
-        {"role": "system", "content": sys_instruction},
-    ]
-
-    for idx in range(len(examples_multimodal_data)):
-        ex_df_txt = examples_multimodal_data[idx]
-        ex_ans = examples_answer[idx]
-        cot = cot_reasoning[idx]
-        messages.extend(
-            [
-                user_message_formatter(ex_df_txt, f"Instruction: {user_instruction}"),
-                {
-                    "role": "assistant",
-                    "content": f"Reasoning:\n{cot}\n\nAnswer: {ex_ans}",
-                },
-            ]
-        )
-
-    messages.append(user_message_formatter(multimodal_data, f"Instruction: {user_instruction}"))
-    return messages
-
-
-def map_formatter_zs_cot(
-    multimodal_data: dict[str, Any],
-    user_instruction: str,
-) -> list[dict[str, str]]:
-    sys_instruction = (
-        "The user will provide an instruction and some relevant context.\n"
-        "Your job is to answer the user's instruction given the context."
-        'First give your reasoning. Then you MUST end your output with "Answer: your answer"'
-    )
-    messages = [
-        {"role": "system", "content": sys_instruction},
-    ]
-
-    messages.append(user_message_formatter(multimodal_data, f"Instruction: {user_instruction}"))
-    return messages
-
-
 def map_formatter(
     multimodal_data: dict[str, Any],
     user_instruction: str,
@@ -201,30 +150,57 @@ def map_formatter(
     examples_answer: list[str] | None = None,
     cot_reasoning: list[str] | None = None,
     strategy: str | None = None,
+    reasoning_instructions: str = "",
 ) -> list[dict[str, str]]:
-    if cot_reasoning:
-        assert examples_multimodal_data is not None and examples_answer is not None
-        return map_formatter_cot(
-            multimodal_data, user_instruction, examples_multimodal_data, examples_answer, cot_reasoning
+    """
+    Creates a map formatter with chain-of-thought reasoning.
+    Supports both few-shot CoT (with examples) and zero-shot CoT.
+    """
+    sys_instruction = """The user will provide an instruction and some relevant context.
+    Your job is to answer the user's instruction given the context.
+    """
+
+    if strategy == "cot":
+        sys_instruction += cot_prompt_formatter(
+            reasoning_instructions=reasoning_instructions,
         )
-    elif strategy == "zs-cot":
-        return map_formatter_zs_cot(multimodal_data, user_instruction)
+    else:
+        sys_instruction += non_cot_prompt_formatter()
 
-    sys_instruction = (
-        "The user will provide an instruction and some relevant context.\n"
-        "Your job is to answer the user's instruction given the context."
-    )
     messages = [
         {"role": "system", "content": sys_instruction},
     ]
 
     if examples_multimodal_data:
         assert examples_answer is not None
-        for ex_df_txt, ex_ans in zip(examples_multimodal_data, examples_answer):
+        assert isinstance(examples_multimodal_data, list) and isinstance(examples_answer, list)
+        assert len(examples_multimodal_data) == len(examples_answer)
+
+        if cot_reasoning:
+            # If CoT reasoning examples are provided, use them
+            assert isinstance(cot_reasoning, list)
+            assert len(examples_multimodal_data) == len(examples_answer) == len(cot_reasoning)
+
+        for idx in range(len(examples_multimodal_data)):
+            ex_df_txt = examples_multimodal_data[idx]
+            ex_ans = examples_answer[idx]
+            content = ""
+
+            # If CoT reasoning is provided, use it. Otherwise, supply a default reasoning
+            if cot_reasoning:
+                content = cot_formatter(cot_reasoning[idx], str(ex_ans))
+            elif strategy == "cot":
+                content = cot_formatter("Reasoning omitted", str(ex_ans))
+            else:
+                content = answer_only_formatter(str(ex_ans))
+
             messages.extend(
                 [
                     user_message_formatter(ex_df_txt, f"Instruction: {user_instruction}"),
-                    {"role": "assistant", "content": str(ex_ans)},
+                    {
+                        "role": "assistant",
+                        "content": content,
+                    },
                 ]
             )
 

From 11e24f9a7534deb6d4523f62fefa77bb902a008e Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 25 Feb 2025 19:44:00 -0800
Subject: [PATCH 02/12] can use default reasoning if none is provided

---
 lotus/sem_ops/sem_map.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/lotus/sem_ops/sem_map.py b/lotus/sem_ops/sem_map.py
index 508f4a82..4a8a8071 100644
--- a/lotus/sem_ops/sem_map.py
+++ b/lotus/sem_ops/sem_map.py
@@ -153,8 +153,7 @@ def __call__(
             examples_multimodal_data = task_instructions.df2multimodal_info(examples, col_li)
             examples_answers = examples["Answer"].tolist()
 
-            if strategy == "cot":
-                return_explanations = True
+            if strategy == "cot" and "Reasoning" in examples.columns:
                 cot_reasoning = examples["Reasoning"].tolist()
 
         output = sem_map(

From 258089d6ab5dd1ea7998a9bc8eebfbfa43850b5b Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 25 Feb 2025 20:13:16 -0800
Subject: [PATCH 03/12] added tests

---
 .github/tests/lm_tests.py | 178 +++++++++++++++++++++++++++++++++++++-
 1 file changed, 177 insertions(+), 1 deletion(-)

diff --git a/.github/tests/lm_tests.py b/.github/tests/lm_tests.py
index 78656047..c06dbc91 100644
--- a/.github/tests/lm_tests.py
+++ b/.github/tests/lm_tests.py
@@ -313,6 +313,182 @@ def test_filter_operation_cot_fewshot_no_reasoning(setup_models, model):
     assert filtered_df.equals(expected_df)
 
 
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_filter_operation_cot(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Text": [
+            "I had two apples, then I gave away one",
+            "My friend gave me an apple",
+            "I gave away both of my apples",
+            "I gave away my apple, then a friend gave me his apple, then I threw my apple away",
+        ]
+    }
+    df = pd.DataFrame(data)
+    user_instruction = "{Text} I have at least one apple"
+    filtered_df = df.sem_filter(user_instruction, strategy="cot")
+    expected_df = pd.DataFrame({"Text": ["I had two apples, then I gave away one", "My friend gave me an apple"]})
+    assert filtered_df.equals(expected_df)
+
+
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_filter_operation_cot_fewshot(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Sequence": [
+            "Five, Four, Three",
+            "A, B, C",
+            "Pond, Lake, Ocean",
+        ]
+    }
+    df = pd.DataFrame(data)
+    examples = {
+        "Sequence": ["1, 2, 3", "penny, nickel, dime, quarter", "villiage, town, city"],
+        "Answer": [True, True, True],
+        "Reasoning": [
+            "1, 2, 3 is an increasing sequence of numbers",
+            "penny, nickel, dime, quarter is an increasing sequence of coins",
+            "villiage, town, city is an increasing sequence of settlements",
+        ],
+    }
+    examples_df = pd.DataFrame(examples)
+
+    user_instruction = "{Sequence} is increasing"
+    filtered_df = df.sem_filter(
+        user_instruction,
+        strategy="cot",
+        examples=examples_df,
+        additional_cot_instructions="Assume the most typical or logical case.",
+    )
+    expected_df = pd.DataFrame(
+        {
+            "Sequence": [
+                "A, B, C",
+                "Pond, Lake, Ocean",
+            ]
+        },
+        index=[1, 2],
+    )
+    assert filtered_df.equals(expected_df)
+
+
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_filter_operation_cot_fewshot_no_reasoning(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Sequence": [
+            "Five, Four, Three",
+            "A, B, C",
+            "Pond, Lake, Ocean",
+        ]
+    }
+    df = pd.DataFrame(data)
+    examples = {
+        "Sequence": ["1, 2, 3", "penny, nickel, dime, quarter", "villiage, town, city"],
+        "Answer": [True, True, True],
+    }
+    examples_df = pd.DataFrame(examples)
+
+    user_instruction = "{Sequence} is increasing"
+    filtered_df = df.sem_filter(user_instruction, strategy="cot", examples=examples_df)
+    expected_df = pd.DataFrame(
+        {
+            "Sequence": [
+                "A, B, C",
+                "Pond, Lake, Ocean",
+            ]
+        },
+        index=[1, 2],
+    )
+    assert filtered_df.equals(expected_df)
+
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_map_operation_cot(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Sequence": [
+            "Alpha, Bravo, Charlie",
+            "One, Two, Three",
+            "Triangle, Square, Pentagon",
+        ]
+    }
+    df = pd.DataFrame(data)
+    user_instruction = "What should be the next item in the sequence: {Sequence}"
+    mapped_df = df.sem_map(user_instruction, strategy="cot")
+    expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
+    assert mapped_df["_map"].equals(expected_df["_map"])
+
+
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_map_operation_cot_fewshot(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Sequence": [
+            "Alpha, Bravo, Charlie",
+            "One, Two, Three",
+            "Triangle, Square, Pentagon",
+        ]
+    }
+    df = pd.DataFrame(data)
+    examples = {
+        "Sequence": ["A, B, C", "Kindergarten, First Grade, Second Grade"],
+        "Answer": ["D", "Third Grade"],
+        "Reasoning": [
+            "D is the next letter in the alphabet after C",
+            "Third Grade is the next grade after Second Grade",
+        ],
+    }
+    examples_df = pd.DataFrame(examples)
+    user_instruction = "What should be the next item in the sequence: {Sequence}"
+    mapped_df = df.sem_map(user_instruction, strategy="cot", examples=examples_df)
+    expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
+    assert mapped_df["_map"].equals(expected_df["_map"])
+
+
+@pytest.mark.parametrize("model", get_enabled("gpt-4o-mini", "ollama/llama3.1"))
+def test_map_operation_cot_fewshot_no_reasoning(setup_models, model):
+    lm = setup_models[model]
+    lotus.settings.configure(lm=lm)
+
+    # Test filter operation on an easy dataframe
+    data = {
+        "Sequence": [
+            "Alpha, Bravo, Charlie",
+            "One, Two, Three",
+            "Triangle, Square, Pentagon",
+        ]
+    }
+    df = pd.DataFrame(data)
+    examples = {
+        "Sequence": ["A, B, C", "Kindergarten, First Grade, Second Grade"],
+        "Answer": ["D", "Third Grade"],
+        "Reasoning": [
+            "D is the next letter in the alphabet after C",
+            "Third Grade is the next grade after Second Grade",
+        ],
+    }
+    examples_df = pd.DataFrame(examples)
+    user_instruction = "What should be the next item in the sequence: {Sequence}"
+    mapped_df = df.sem_map(user_instruction, strategy="cot", examples=examples_df)
+    expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
+    assert mapped_df["_map"].equals(expected_df["_map"])
+
+
 ################################################################################
 # Cascade tests
 ################################################################################
@@ -392,7 +568,7 @@ def test_filter_cascade(setup_models):
 def test_join_cascade(setup_models):
     models = setup_models
     rm = SentenceTransformersRM(model="intfloat/e5-base-v2")
-    vs = FaissVS() 
+    vs = FaissVS()
     lotus.settings.configure(lm=models["gpt-4o-mini"], rm=rm, vs=vs)
 
     data1 = {

From d8644f00bc3bb669cf0ca122c406f10c63ecdc0d Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 25 Feb 2025 20:13:49 -0800
Subject: [PATCH 04/12] fixed test

---
 .github/tests/lm_tests.py | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/.github/tests/lm_tests.py b/.github/tests/lm_tests.py
index c06dbc91..690c0bdc 100644
--- a/.github/tests/lm_tests.py
+++ b/.github/tests/lm_tests.py
@@ -477,10 +477,6 @@ def test_map_operation_cot_fewshot_no_reasoning(setup_models, model):
     examples = {
         "Sequence": ["A, B, C", "Kindergarten, First Grade, Second Grade"],
         "Answer": ["D", "Third Grade"],
-        "Reasoning": [
-            "D is the next letter in the alphabet after C",
-            "Third Grade is the next grade after Second Grade",
-        ],
     }
     examples_df = pd.DataFrame(examples)
     user_instruction = "What should be the next item in the sequence: {Sequence}"

From 5d7514b16c3c741c1cfbc3b3f1d3761cc599ea0d Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 25 Feb 2025 20:21:20 -0800
Subject: [PATCH 05/12] bugfix

---
 lotus/sem_ops/postprocessors.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index ea6eb4b7..27c2d69b 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -12,15 +12,15 @@ def cot_postprocessor(llm_answers: list[str]):
     outputs: list[str | None] = []
     explanations: list[str | None] = []
     for llm_answer in llm_answers:
-        reasoning_idx = llm_answer.find("Reasoning:\n")
+        reasoning_idx = llm_answer.find("Reasoning: \n")
         if reasoning_idx == -1:
             reasoning_idx = 0
         else:
             reasoning_idx += len("Reasoning:\n")
 
-        answer_idx = llm_answer.find("Answer:")
+        answer_idx = llm_answer.find("Answer: ")
         reasoning = llm_answer[reasoning_idx:answer_idx].rstrip("\n").lstrip("\n")
-        answer = llm_answer[answer_idx + len("Answer:") :]
+        answer = llm_answer[answer_idx + len("Answer: ") :]
 
         explanations.append(reasoning)
         outputs.append(answer)

From 76a28d59c52277973bc5a9a0d56142c0357e87e2 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 09:56:12 -0700
Subject: [PATCH 06/12] handle case where LLM doesn't follow formatting for
 answer

---
 lotus/sem_ops/postprocessors.py | 10 ++++++++--
 1 file changed, 8 insertions(+), 2 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index 27c2d69b..705a92f8 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -12,15 +12,21 @@ def cot_postprocessor(llm_answers: list[str]):
     outputs: list[str | None] = []
     explanations: list[str | None] = []
     for llm_answer in llm_answers:
-        reasoning_idx = llm_answer.find("Reasoning: \n")
+        reasoning_idx = llm_answer.find("Reasoning:\n")
         if reasoning_idx == -1:
             reasoning_idx = 0
         else:
             reasoning_idx += len("Reasoning:\n")
 
         answer_idx = llm_answer.find("Answer: ")
+        if answer_idx == -1:
+            answer_idx = 0
+        else:
+            answer_idx += len("Answer: ")
+
+
         reasoning = llm_answer[reasoning_idx:answer_idx].rstrip("\n").lstrip("\n")
-        answer = llm_answer[answer_idx + len("Answer: ") :]
+        answer = llm_answer[answer_idx:].rstrip("\n").lstrip("\n")
 
         explanations.append(reasoning)
         outputs.append(answer)

From e811f372b549a604542f93cb10864cf19d21af03 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 09:58:10 -0700
Subject: [PATCH 07/12] fix logging

---
 lotus/sem_ops/postprocessors.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index 705a92f8..0553e90f 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -9,8 +9,8 @@
 
 
 def cot_postprocessor(llm_answers: list[str]):
-    outputs: list[str | None] = []
-    explanations: list[str | None] = []
+    outputs: list[str] = []
+    explanations: list[str] = []
     for llm_answer in llm_answers:
         reasoning_idx = llm_answer.find("Reasoning:\n")
         if reasoning_idx == -1:
@@ -49,7 +49,7 @@ def map_postprocess(llm_answers: list[str], default: str = "") -> SemanticMapPos
 
     for i, output in enumerate(outputs):
         if output is None:
-            lotus.logger.info(f"\t Failed to parse {output}: defaulting to {default}")
+            lotus.logger.info(f"\t Failed to parse output {i}: defaulting to {default}")
             outputs[i] = default
 
     return SemanticMapPostprocessOutput(raw_outputs=llm_answers, outputs=outputs, explanations=explanations)

From 04fda6ac81566564c8effc8ae8c43b1c8cb82741 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 10:04:06 -0700
Subject: [PATCH 08/12] fix logging

---
 lotus/sem_ops/postprocessors.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index 0553e90f..dc40de00 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -49,7 +49,7 @@ def map_postprocess(llm_answers: list[str], default: str = "") -> SemanticMapPos
 
     for i, output in enumerate(outputs):
         if output is None:
-            lotus.logger.info(f"\t Failed to parse output {i}: defaulting to {default}")
+            lotus.logger.info(f"\t Failed to parse {llm_answers[i]}: defaulting to {default}")
             outputs[i] = default
 
     return SemanticMapPostprocessOutput(raw_outputs=llm_answers, outputs=outputs, explanations=explanations)

From e7ee0f1375dd93a12dd552c3cf9fd17a88ada0a1 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 10:13:59 -0700
Subject: [PATCH 09/12] fix docstring

---
 lotus/sem_ops/sem_map.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/lotus/sem_ops/sem_map.py b/lotus/sem_ops/sem_map.py
index 4a8a8071..fbff4b36 100644
--- a/lotus/sem_ops/sem_map.py
+++ b/lotus/sem_ops/sem_map.py
@@ -36,6 +36,11 @@ def sem_map(
         examples_multimodal_data (list[dict[str, Any]] | None): The text for examples. Defaults to None.
         examples_answers (list[str] | None): The answers for examples. Defaults to None.
         cot_reasoning (list[str] | None): The reasoning for CoT. Defaults to None.
+        additional_cot_instructions (str): Additional instructions for the CoT. Defaults to "".
+        strategy (str | None): The reasoning strategy. Defaults to None.
+        safe_mode (bool): Whether to use safe mode. Defaults to False.
+        progress_bar_desc (str): The description for the progress bar. Defaults to "Mapping".
+        default (str): The default value to use if we fail to parse the answer.
 
     Returns:
         SemanticMapOutput: The outputs, raw outputs, and explanations.
@@ -126,6 +131,7 @@ def __call__(
             safe_mode (bool): Whether to use safe mode. Defaults to False.
             progress_bar_desc (str): The description for the progress bar. Defaults to "Mapping".
             additional_cot_instructions (str): Additional instructions for the CoT. Defaults to "".
+            default (str): The default value to use if we fail to parse the answer.
 
         Returns:
             pd.DataFrame: The dataframe with the new mapped columns.

From 13adc980814cd5ac57b8698440098be3678a62d8 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 11:04:45 -0700
Subject: [PATCH 10/12] bug fixes post merge

---
 lotus/sem_ops/postprocessors.py      | 5 +++--
 lotus/templates/task_instructions.py | 4 ++--
 2 files changed, 5 insertions(+), 4 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index 97dcb5ef..11aa4cb5 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -11,7 +11,7 @@
 
 def cot_postprocessor(llm_answers: list[str]):
     outputs: list[str] = []
-    explanations: list[str] = ]
+    explanations: list[str] = []
     for llm_answer in llm_answers:
         reasoning_idx = llm_answer.find("Reasoning:\n")
         if reasoning_idx == -1:
@@ -31,7 +31,7 @@ def cot_postprocessor(llm_answers: list[str]):
 
         explanations.append(reasoning)
         outputs.append(answer)
-
+    
     return outputs, explanations
 
 
@@ -127,6 +127,7 @@ def map_postprocess(llm_answers: list[str], model: lotus.models.LM, cot_reasonin
     if cot_reasoning:
         postprocessor = get_cot_postprocessor(model)
         outputs, explanations = postprocessor(llm_answers)
+        outputs = [output if output is not None else default for output in outputs]
     else:
         outputs = llm_answers
         explanations = [None] * len(llm_answers)
diff --git a/lotus/templates/task_instructions.py b/lotus/templates/task_instructions.py
index 05e12c21..30d4f156 100644
--- a/lotus/templates/task_instructions.py
+++ b/lotus/templates/task_instructions.py
@@ -171,7 +171,7 @@ def map_formatter(
     Your job is to answer the user's instruction given the context.
     """
 
-    if strategy == "cot":
+    if strategy in [ReasoningStrategy.COT, ReasoningStrategy.ZS_COT]:
         sys_instruction += cot_prompt_formatter(
             reasoning_instructions=reasoning_instructions,
         )
@@ -200,7 +200,7 @@ def map_formatter(
             # If CoT reasoning is provided, use it. Otherwise, supply a default reasoning
             if cot_reasoning:
                 content = cot_formatter(cot_reasoning[idx], str(ex_ans))
-            elif strategy == "cot":
+            elif strategy in [ReasoningStrategy.COT, ReasoningStrategy.ZS_COT]:
                 content = cot_formatter("Reasoning omitted", str(ex_ans))
             else:
                 content = answer_only_formatter(str(ex_ans))

From 4ac1fe61e80f0ccbb67e65e9d47139f768087f53 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 11:25:06 -0700
Subject: [PATCH 11/12] fixed tests

---
 .github/tests/lm_tests.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/.github/tests/lm_tests.py b/.github/tests/lm_tests.py
index b79bcf00..0e48970c 100644
--- a/.github/tests/lm_tests.py
+++ b/.github/tests/lm_tests.py
@@ -6,7 +6,7 @@
 
 import lotus
 from lotus.models import LM, SentenceTransformersRM
-from lotus.types import CascadeArgs
+from lotus.types import CascadeArgs, ReasoningStrategy
 from lotus.vector_store import FaissVS
 
 ################################################################################
@@ -269,7 +269,7 @@ def test_filter_operation_cot(setup_models, model):
     }
     df = pd.DataFrame(data)
     user_instruction = "{Text} I have at least one apple"
-    filtered_df = df.sem_filter(user_instruction, strategy="cot")
+    filtered_df = df.sem_filter(user_instruction, strategy=ReasoningStrategy.ZS_COT)
     expected_df = pd.DataFrame({"Text": ["I had two apples, then I gave away one", "My friend gave me an apple"]})
     assert filtered_df.equals(expected_df)
 
@@ -302,7 +302,7 @@ def test_filter_operation_cot_fewshot(setup_models, model):
     user_instruction = "{Sequence} is increasing"
     filtered_df = df.sem_filter(
         user_instruction,
-        strategy="cot",
+        strategy=ReasoningStrategy.COT,
         examples=examples_df,
         additional_cot_instructions="Assume the most typical or logical case.",
     )
@@ -339,7 +339,7 @@ def test_filter_operation_cot_fewshot_no_reasoning(setup_models, model):
     examples_df = pd.DataFrame(examples)
 
     user_instruction = "{Sequence} is increasing"
-    filtered_df = df.sem_filter(user_instruction, strategy="cot", examples=examples_df)
+    filtered_df = df.sem_filter(user_instruction, strategy=ReasoningStrategy.ZS_COT, examples=examples_df)
     expected_df = pd.DataFrame(
         {
             "Sequence": [
@@ -368,7 +368,7 @@ def test_filter_operation_cot(setup_models, model):
     }
     df = pd.DataFrame(data)
     user_instruction = "{Text} I have at least one apple"
-    filtered_df = df.sem_filter(user_instruction, strategy="cot")
+    filtered_df = df.sem_filter(user_instruction, strategy=ReasoningStrategy.ZS_COT)
     expected_df = pd.DataFrame({"Text": ["I had two apples, then I gave away one", "My friend gave me an apple"]})
     assert filtered_df.equals(expected_df)
 
@@ -401,7 +401,7 @@ def test_filter_operation_cot_fewshot(setup_models, model):
     user_instruction = "{Sequence} is increasing"
     filtered_df = df.sem_filter(
         user_instruction,
-        strategy="cot",
+        strategy=ReasoningStrategy.COT,
         examples=examples_df,
         additional_cot_instructions="Assume the most typical or logical case.",
     )
@@ -438,7 +438,7 @@ def test_filter_operation_cot_fewshot_no_reasoning(setup_models, model):
     examples_df = pd.DataFrame(examples)
 
     user_instruction = "{Sequence} is increasing"
-    filtered_df = df.sem_filter(user_instruction, strategy="cot", examples=examples_df)
+    filtered_df = df.sem_filter(user_instruction, strategy=ReasoningStrategy.ZS_COT, examples=examples_df)
     expected_df = pd.DataFrame(
         {
             "Sequence": [
@@ -465,7 +465,7 @@ def test_map_operation_cot(setup_models, model):
     }
     df = pd.DataFrame(data)
     user_instruction = "What should be the next item in the sequence: {Sequence}"
-    mapped_df = df.sem_map(user_instruction, strategy="cot")
+    mapped_df = df.sem_map(user_instruction, strategy=ReasoningStrategy.ZS_COT)
     expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
     assert mapped_df["_map"].equals(expected_df["_map"])
 
@@ -494,7 +494,7 @@ def test_map_operation_cot_fewshot(setup_models, model):
     }
     examples_df = pd.DataFrame(examples)
     user_instruction = "What should be the next item in the sequence: {Sequence}"
-    mapped_df = df.sem_map(user_instruction, strategy="cot", examples=examples_df)
+    mapped_df = df.sem_map(user_instruction, strategy=ReasoningStrategy.COT, examples=examples_df)
     expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
     assert mapped_df["_map"].equals(expected_df["_map"])
 
@@ -519,7 +519,7 @@ def test_map_operation_cot_fewshot_no_reasoning(setup_models, model):
     }
     examples_df = pd.DataFrame(examples)
     user_instruction = "What should be the next item in the sequence: {Sequence}"
-    mapped_df = df.sem_map(user_instruction, strategy="cot", examples=examples_df)
+    mapped_df = df.sem_map(user_instruction, strategy=ReasoningStrategy.ZS_COT, examples=examples_df)
     expected_df = pd.DataFrame({"_map": ["Delta", "Four", "Hexagon"]})
     assert mapped_df["_map"].equals(expected_df["_map"])
 

From 2a9740afb9b1208f0430c2ae3bfe35598310aea2 Mon Sep 17 00:00:00 2001
From: dhruviyer <iyer.dhruv99@gmail.com>
Date: Tue, 15 Apr 2025 11:34:18 -0700
Subject: [PATCH 12/12] fix post processing to remove answer even when COT is
 not used

---
 lotus/sem_ops/postprocessors.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/lotus/sem_ops/postprocessors.py b/lotus/sem_ops/postprocessors.py
index 11aa4cb5..364cf594 100644
--- a/lotus/sem_ops/postprocessors.py
+++ b/lotus/sem_ops/postprocessors.py
@@ -124,13 +124,9 @@ def map_postprocess(llm_answers: list[str], model: lotus.models.LM, cot_reasonin
         SemanticMapPostprocessOutput
     """
 
-    if cot_reasoning:
-        postprocessor = get_cot_postprocessor(model)
-        outputs, explanations = postprocessor(llm_answers)
-        outputs = [output if output is not None else default for output in outputs]
-    else:
-        outputs = llm_answers
-        explanations = [None] * len(llm_answers)
+    postprocessor = get_cot_postprocessor(model)
+    outputs, explanations = postprocessor(llm_answers)
+    outputs = [output if output is not None else default for output in outputs]
 
     return SemanticMapPostprocessOutput(raw_outputs=llm_answers, outputs=outputs, explanations=explanations)