Repository: argilla-io/distilabel

Issue:
[BUG] Error when wrapping the step
**Describe the bug**
A clear and concise description of what the bug is.
When working on this notebook (https://github.com/argilla-io/distilabel/pull/949), I encountered that with the dev branch, the task wasn't working.

Maybe related to this one too: https://github.com/argilla-io/distilabel/issues/991

**To Reproduce**
Code to reproduce
```python
labels_topic= ["World", "Sports", "Sci/Tech", "Business"]
labels_fact_opinion = ["Fact-based", "Opinion-based"]

task_templates = [
    "Determine the news article as {}",
    "Classify news article as {}",
    "Identify the news article as {}",
    "Categorize the news article as {}",
    "Label the news article using {}",
    "Annotate the news article based on {}",
    "Determine the theme of a news article from {}",
    "Recognize the topic of the news article as {}",
]

classification_tasks = [
    {"task": action.format(" or ".join(random.sample(labels_topic, 2)))}
    for action in task_templates for _ in range(4)
] + [
    {"task": action.format(" or ".join(random.sample(labels_fact_opinion, 2)))}
    for action in task_templates
]

difficulties = ["college", "high school", "PhD"]
clarity = ["clear", "understandable with some effort", "ambiguous"]

with Pipeline("texcat-generation-pipeline") as pipeline:

    tasks_generator = LoadDataFromDicts(data=classification_tasks)

    generate_data = []
    for difficulty in difficulties:
        for clarity_level in clarity:
            task = GenerateTextClassificationData(
                language="English",
                difficulty=difficulty,
                clarity=clarity_level,
                num_generations=2,
                llm=InferenceEndpointsLLM(
                    model_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
                    tokenizer_id="meta-llama/Meta-Llama-3.1-8B-Instruct",
                    generation_kwargs={"max_new_tokens": 512, "temperature": 0.7},
                ),
                input_batch_size=5,
            )
            generate_data.append(task)

    for task in generate_data:
        tasks_generator.connect(task)
```

**Expected behaviour**
A clear and concise description of what you expected to happen.

**Screenshots**
If applicable, add screenshots to help explain your problem.
![Screenshot 2024-09-25 at 12 20 10](https://github.com/user-attachments/assets/37f3e555-bd08-48a2-ade5-e8c02f118819)


**Desktop (please complete the following information):**
 - Package version: 1.4dev
 - Python version: 3.11.4

**Additional context**
Add any other context about the problem here.
Note: Already shared on Slack, just realized I didn't create an issue

Hints from issue comments:
@plaguss: Hi @sdiazlor, can you share the `classification_tasks` data to run the example?@sdiazlor: @plaguss thanks, I updated the example code. With the 1.3.2 works well.@plaguss: Can you test with the fix and see if it works?@sdiazlor: It works now, thanks for tackling this!

Solution patch:
diff --git a/src/distilabel/steps/base.py b/src/distilabel/steps/base.py
index c467005e0..b98c0e827 100644
--- a/src/distilabel/steps/base.py
+++ b/src/distilabel/steps/base.py
@@ -667,10 +667,19 @@ class Step(_Step, ABC):
         )
 
         for output_rows in generator:
-            yield [
-                self._apply_mappings_and_restore_overriden(row, overriden_inputs[i])
-                for i, row in enumerate(output_rows)
-            ]
+            restored = []
+            for i, row in enumerate(output_rows):
+                # Correct the index here because we don't know the num_generations from the llm
+                # ahead of time. For example, if we have `len(overriden_inputs)==5` and `len(row)==10`,
+                # from `num_generations==2` and `group_generations=False` in the LLM:
+                # The loop will use indices 0, 1, 2, 3, 4, 0, 1, 2, 3, 4
+                ntimes_i = i % len(overriden_inputs)
+                restored.append(
+                    self._apply_mappings_and_restore_overriden(
+                        row, overriden_inputs[ntimes_i]
+                    )
+                )
+            yield restored
 
     def _apply_input_mappings(
         self, inputs: Tuple[List[Dict[str, Any]], ...]


Test patch:
diff --git a/tests/unit/steps/tasks/test_base.py b/tests/unit/steps/tasks/test_base.py
index c56991d47..29341052f 100644
--- a/tests/unit/steps/tasks/test_base.py
+++ b/tests/unit/steps/tasks/test_base.py
@@ -414,6 +414,94 @@ class TestTask:
         result = next(task.process(input))
         assert result == expected
 
+    def test_process_overriding_inputs(self) -> None:
+        llm = DummyAsyncLLM()
+        task = DummyTask(
+            name="task",
+            llm=llm,
+            group_generations=False,
+            num_generations=3,
+            input_mappings={"instruction": "instruction_2"},
+        )
+
+        result = next(
+            task.process_applying_mappings(
+                [
+                    {
+                        "instruction": "instruction that won't be used but overriden by input mapping",
+                        "instruction_2": "instruction that will be used as input",
+                        "additional_info": "info",
+                    }
+                ]
+            )
+        )
+
+        assert result == [
+            {
+                "additional_info": "info",
+                "distilabel_metadata": {
+                    "raw_input_task": [
+                        {
+                            "content": "",
+                            "role": "system",
+                        },
+                        {
+                            "content": "instruction that will be used as input",
+                            "role": "user",
+                        },
+                    ],
+                    "raw_output_task": "output",
+                },
+                "info_from_input": "info",
+                "instruction": "instruction that won't be used but overriden by input mapping",
+                "instruction_2": "instruction that will be used as input",
+                "model_name": "test",
+                "output": "output",
+            },
+            {
+                "additional_info": "info",
+                "distilabel_metadata": {
+                    "raw_input_task": [
+                        {
+                            "content": "",
+                            "role": "system",
+                        },
+                        {
+                            "content": "instruction that will be used as input",
+                            "role": "user",
+                        },
+                    ],
+                    "raw_output_task": "output",
+                },
+                "info_from_input": "info",
+                "instruction": "instruction that won't be used but overriden by input mapping",
+                "instruction_2": "instruction that will be used as input",
+                "model_name": "test",
+                "output": "output",
+            },
+            {
+                "additional_info": "info",
+                "distilabel_metadata": {
+                    "raw_input_task": [
+                        {
+                            "content": "",
+                            "role": "system",
+                        },
+                        {
+                            "content": "instruction that will be used as input",
+                            "role": "user",
+                        },
+                    ],
+                    "raw_output_task": "output",
+                },
+                "info_from_input": "info",
+                "instruction": "instruction that won't be used but overriden by input mapping",
+                "instruction_2": "instruction that will be used as input",
+                "model_name": "test",
+                "output": "output",
+            },
+        ]
+
     def test_process_with_runtime_parameters(self) -> None:
         # 1. Runtime parameters provided
         llm = DummyRuntimeLLM()  # type: ignore