Merge pull request #1 from huggingface/generate-dpo-dataset

data/download_data.sh

@@ -0,0 +1 @@
+wget https://huggingface.co/datasets/feel-fl/open-human-feedback-chat-en/resolve/main/data/data_d6f0f072-348e-4f61-9a44-26dbd2ccba75.json

data/example_data.json

@@ -0,0 +1,59 @@
+[
+    {
+        "conversation": [
+            {
+                "role": "user",
+                "content": "hello how are you?",
+                "options": "",
+                "rating": 0
+            },
+            {
+                "role": "assistant",
+                "content": "leave me alone you weirdo!",
+                "options": "",
+                "rating": -1
+            }
+        ],
+        "timestamp": "2024-12-10T15:35:52.363635",
+        "session_id": "9c5b367d-12c2-4ae0-a868-e2e783e50935",
+        "conversation_id": "870fac58-2b2c-45ac-93f7-7cd8a43644be"
+    },
+    {
+        "conversation": [
+            {
+                "role": "user",
+                "content": "hello",
+                "options": "",
+                "rating": 0
+            },
+            {
+                "role": "assistant",
+                "content": "Hello! How can I assist you today? If you have any questions or just want to chat, feel free \ud83d\ude0a.",
+                "options": "",
+                "rating": 1
+            }
+        ],
+        "timestamp": "2024-12-10T15:35:52.363635",
+        "session_id": "9c5b367d-12c2-4ae0-a868-e2e783e50935",
+        "conversation_id": "870fac58-2b2c-45ac-93f7-7cd8a43644be"
+    },
+    {
+        "conversation": [
+            {
+                "role": "user",
+                "content": "hello",
+                "options": "",
+                "rating": 0
+            },
+            {
+                "role": "assistant",
+                "content": "Hello! How can I assist you today? If you have any questions or just want to chat, feel free \ud83d\ude0a.",
+                "options": "",
+                "rating": 1
+            }
+        ],
+        "timestamp": "2024-12-10T15:35:52.363635",
+        "session_id": "9c5b367d-12c2-4ae0-a868-e2e783e50935",
+        "conversation_id": "870fac58-2b2c-45ac-93f7-7cd8a43644be"
+    }
+]

data/generate_dpo.py

@@ -0,0 +1,167 @@
+import json
+from typing import TYPE_CHECKING, List, Literal, Union
+
+from datasets import Dataset, concatenate_datasets
+from distilabel.llms.huggingface import InferenceEndpointsLLM
+from distilabel.pipeline import Pipeline
+from distilabel.steps import CombineOutputs, GeneratorStep, KeepColumns, Step, StepInput
+from distilabel.steps.tasks import TextGeneration
+from typing_extensions import override
+
+CHOSEN_TEMPLATE = """
+You are provide with a conversation between a human and an AI assistant. 
+The final message is of poor quality positively. Your task is to regenerate one of high quality.
+{% for message in conversation %}
+{{ message["role"] }}: {{ message["content"] }}
+{% endfor %}
+High quality response:
+""".rstrip()
+
+CHOSEN_SYSTEM_PROMPT = "You are a helpful AI assistant. Your task is to generate high quality response when other assistants created a poor quality response."
+
+REJECT_TEMPLATE = """
+You are provide with a conversation between a human and an AI assistant.
+The final message is of high quality positively. Your task is to regenerate one of poor quality.
+{% for message in conversation %}
+{{ message["role"] }}: {{ message["content"] }}
+{% endfor %}
+Poor quality response:
+""".rstrip()
+
+REJECT_SYSTEM_PROMPT = "You are a helpful AI assistant. Your task is to generate a poor quality response when other assistants created a high quality response."
+
+
+class FilterConversationRatings(Step):
+    """Filters conversations based on the rating of the last message."""
+
+    target_column: Union[Literal["chosen"], Literal["rejected"]]
+    batch_size: int = 5
+
+    @override
+    def process(self, dataset: StepInput) -> "GeneratorStepOutput":
+
+        column_rating_map = {
+            "chosen": 1,
+            "rejected": -1,
+        }
+
+        target_rating = column_rating_map[self.target_column]
+
+        for batch_start in range(0, len(dataset), self.batch_size):
+            batch = dataset[batch_start : batch_start + self.batch_size]
+            filtered_batch = []
+            for conversation in batch:
+                for row in batch:
+                    _conversation = row["conversation"]
+                    conversation = None
+                    for idx, message in enumerate(_conversation, 1):
+                        if not isinstance(message["rating"], int):
+                            continue
+                        if message["rating"] == target_rating:
+                            conversation = _conversation[:idx]
+                            break
+                    if conversation:
+                        filtered_batch.append({"conversation": conversation})
+            yield filtered_batch
+
+    @property
+    def outputs(self) -> "StepColumns":
+        return ["conversation"]
+
+
+class AppendToConversationStep(Step):
+    """Appends a generated message to a conversation."""
+
+    @property
+    def inputs(self) -> "StepColumns":
+        return ["generation", "conversation"]
+
+    @property
+    def outputs(self) -> "StepColumns":
+        return ["generated_conversation", "conversation"]
+
+    def process(self, inputs: StepInput) -> "StepOutput":
+
+        for input in inputs:
+            if not input["generation"]:
+                continue
+            if not input["conversation"]:
+                continue
+            input["generated_conversation"] = [
+                {"role": message["role"], "content": message["content"]}
+                for message in input["conversation"][:-1]
+            ] + [{"role": "assistant", "content": input["generation"]}]
+            input["conversation"] = [
+                {"role": message["role"], "content": message["content"]}
+                for message in input["conversation"]
+            ]
+        yield inputs
+
+
+with Pipeline(
+    name="conversation_rejection",
+    description="Generate a chosen response to a rejected conversation.",
+) as rejection_pipeline:
+
+    rejected_dataset = FilterConversationRatings(target_column="rejected")
+
+    chosen_text_gen = TextGeneration(
+        llm=InferenceEndpointsLLM(
+            model_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
+        ),
+        system_prompt=CHOSEN_SYSTEM_PROMPT,
+        template=CHOSEN_TEMPLATE,
+        columns=["conversation"],
+    )
+
+    append_chosen = AppendToConversationStep(
+        output_mappings={
+            "generated_conversation": "chosen",
+            "conversation": "rejected",
+        },
+    )
+
+    keep_columns = KeepColumns(
+        columns=["chosen", "rejected"],
+    )
+
+    rejected_dataset >> chosen_text_gen >> append_chosen >> keep_columns
+
+with Pipeline(
+    name="conversation_chosen",
+    description="Generate a rejected response to a chosen conversation.",
+) as chosen_pipeline:
+
+    chosen_dataset = FilterConversationRatings(target_column="chosen")
+
+    rejected_text_gen = TextGeneration(
+        llm=InferenceEndpointsLLM(
+            model_id="meta-llama/Meta-Llama-3.1-70B-Instruct",
+        ),
+        system_prompt=REJECT_SYSTEM_PROMPT,
+        template=REJECT_TEMPLATE,
+        columns=["conversation"],
+    )
+    append_rejected = AppendToConversationStep(
+        output_mappings={
+            "generated_conversation": "rejected",
+            "conversation": "chosen",
+        },
+    )
+    keep_columns = KeepColumns(
+        columns=["chosen", "rejected"],
+    )
+    chosen_dataset >> rejected_text_gen >> append_rejected >> keep_columns
+
+if __name__ == "__main__":
+
+    dataset_path = "example_data.json"
+    data = json.load(open(dataset_path))
+
+    dataset = Dataset.from_list(data)
+    rejected_dataset = rejection_pipeline.run(dataset=dataset, use_cache=False)
+    chosen_dataset = chosen_pipeline.run(dataset=dataset, use_cache=False)
+
+    dataset = concatenate_datasets(
+        dsets=[rejected_dataset["default"]["train"], chosen_dataset["default"]["train"]]
+    )

pyproject.toml

@@ -6,6 +6,8 @@ readme = "README.md"
 requires-python = ">=3.11"
 dependencies = [
     "datasets>=3.1.0",
+    "distilabel>=1.4.1",
+    "ipykernel>=6.29.5",
 ]
 
 [dependency-groups]