Repository: openai/openai-python

Issue:
KeyError when removing long examples after removing duplicate rows
Error: 
```
openai tools fine_tunes.prepare_data -f training_data_2022-09-14.jsonl
Analyzing...

- Your file contains 2446 prompt-completion pairs
Based on the analysis we will perform the following actions:
- [Recommended] Remove 1155 duplicate rows [Y/n]: y
- [Recommended] Remove 49 long examples [Y/n]: y
Traceback (most recent call last):
  File "/Users/ser/project/project-venv/bin/openai", line 8, in <module>
    sys.exit(main())
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/openai/_openai_scripts.py", line 63, in main
    args.func(args)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/openai/cli.py", line 531, in prepare_data
    apply_validators(
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/openai/validators.py", line 851, in apply_validators
    df, optional_applied = apply_optional_remediation(
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/openai/validators.py", line 578, in apply_optional_remediation
    df = remediation.optional_fn(df)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/openai/validators.py", line 171, in optional_fn
    return x.drop(long_indexes)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/pandas/util/_decorators.py", line 311, in wrapper
    return func(*args, **kwargs)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/pandas/core/frame.py", line 4957, in drop
    return super().drop(
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/pandas/core/generic.py", line 4267, in drop
    obj = obj._drop_axis(labels, axis, level=level, errors=errors)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/pandas/core/generic.py", line 4311, in _drop_axis
    new_axis = axis.drop(labels, errors=errors)
  File "/Users/ser/project/project-venv/lib/python3.10/site-packages/pandas/core/indexes/base.py", line 6661, in drop
    raise KeyError(f"{list(labels[mask])} not found in axis")
KeyError: '[330, 352, 377, 378, 422, 424, 435, 1172, 1194, 1219, 1220, 1264, 1266, 1277, 1468, 1498, 1549, 1641, 1648, 1714, 1741, 1816, 1859, 1984] not found in axis'
```

I believe that since the duplicate rows were removed, many of the long examples are missing, throwing this error. And thus I end up needing to apply the first recommendation and not the second one, and then use the resulting file to apply the second recommendation.

It would be great to be able to apply both changes to the same file.


Discussion:
Hi @serinamarie! Thanks for the issue!

Yep, that sure looks like what's happening. We should be able to take a look at this soon and dive into why this is happening (or if you want to take a shot at it, I certainly won't stop you)
Sounds good, I believe I know what the issue is so I can get to it this weekend 👍 

Resolving patch:
diff --git a/openai/validators.py b/openai/validators.py
--- a/openai/validators.py
+++ b/openai/validators.py
@@ -158,17 +158,24 @@ def long_examples_validator(df):
 
     ft_type = infer_task_type(df)
     if ft_type != "open-ended generation":
-        long_examples = df.apply(
-            lambda x: len(x.prompt) + len(x.completion) > 10000, axis=1
-        )
-        long_indexes = df.reset_index().index[long_examples].tolist()
+        def get_long_indexes(d):
+            long_examples = d.apply(
+                lambda x: len(x.prompt) + len(x.completion) > 10000, axis=1
+            )
+            return d.reset_index().index[long_examples].tolist()
+
+        long_indexes = get_long_indexes(df)
 
         if len(long_indexes) > 0:
             immediate_msg = f"\n- There are {len(long_indexes)} examples that are very long. These are rows: {long_indexes}\nFor conditional generation, and for classification the examples shouldn't be longer than 2048 tokens."
             optional_msg = f"Remove {len(long_indexes)} long examples"
 
             def optional_fn(x):
-                return x.drop(long_indexes)
+                
+                long_indexes_to_drop = get_long_indexes(x)
+                if long_indexes != long_indexes_to_drop:
+                    sys.stdout.write(f"The indices of the long examples has changed as a result of a previously applied recommendation.\nThe {len(long_indexes_to_drop)} long examples to be dropped are now at the following indices: {long_indexes_to_drop}\n")
+                return x.drop(long_indexes_to_drop)
 
     return Remediation(
         name="long_examples",


Test patch:
diff --git a/openai/tests/test_long_examples_validator.py b/openai/tests/test_long_examples_validator.py
new file mode 100644
--- /dev/null
+++ b/openai/tests/test_long_examples_validator.py
@@ -0,0 +1,48 @@
+import json
+import subprocess
+from tempfile import NamedTemporaryFile
+
+
+def test_long_examples_validator() -> None:
+
+    """
+    Ensures that long_examples_validator() handles previously applied recommendations,
+    namely dropped duplicates, without resulting in a KeyError.
+    """
+
+    # data
+    short_prompt = "a prompt "
+    long_prompt = short_prompt * 500
+
+    short_completion = "a completion "
+    long_completion = short_completion * 500
+
+    # the order of these matters
+    unprepared_training_data = [
+        {"prompt": long_prompt, "completion": long_completion},  # 1 of 2 duplicates
+        {"prompt": short_prompt, "completion": short_completion}, 
+        {"prompt": long_prompt, "completion": long_completion},  # 2 of 2 duplicates
+
+    ]
+
+    with NamedTemporaryFile(suffix="jsonl", mode="w") as training_data:
+        for prompt_completion_row in unprepared_training_data:
+            training_data.write(json.dumps(prompt_completion_row) + "\n")
+            training_data.flush()
+    
+        prepared_data_cmd_output = subprocess.run(
+            [f"openai tools fine_tunes.prepare_data -f {training_data.name}"], 
+            stdout=subprocess.PIPE, 
+            text=True, 
+            input="y\ny\ny\ny\ny",  # apply all recommendations, one at a time
+            stderr=subprocess.PIPE,
+            encoding="utf-8",
+            shell=True
+        )
+
+    # validate data was prepared successfully
+    assert prepared_data_cmd_output.stderr == ""  
+    # validate get_long_indexes() applied during optional_fn() call in long_examples_validator()
+    assert "indices of the long examples has changed" in prepared_data_cmd_output.stdout
+    
+    return prepared_data_cmd_output.stdout
\ No newline at end of file