reafactor: 重构项目

pipeline/processors/dpo.py

@@ -0,0 +1,51 @@
+"""DPO preference learning data processor."""
+from typing import Dict, List, Any, Optional
+
+import torch
+from torch import Tensor
+
+from pipeline.tokenizer import BpeTokenizer
+from pipeline.strategies import PromptStrategy, ChatMLStrategy
+from pipeline.processors.base import BaseProcessor, _encode_with_mask
+from pipeline.processors.factory import ProcessorFactory
+
+
+@ProcessorFactory.register("dpo")
+class DPOProcessor(BaseProcessor):
+    """DPO preference learning data processor.
+
+    Supports custom prompt strategy via constructor parameter.
+    """
+
+    def __init__(
+        self,
+        tokenizer: BpeTokenizer,
+        strategy: Optional[PromptStrategy] = None,
+    ):
+        self.tokenizer = tokenizer
+        self.strategy = strategy or ChatMLStrategy(tokenizer)
+
+    def process(self, input_dict: Dict[str, Any]) -> Dict[str, Tensor]:
+        query_tokens = self.tokenizer.encode(input_dict["query"])
+        chosen_tokens = self.tokenizer.encode(input_dict["chosen"])
+        rejected_tokens = self.tokenizer.encode(input_dict["rejected"])
+
+        prompt = self.strategy.assemble_prompt(query_tokens)
+
+        chosen_t, chosen_m = _encode_with_mask(
+            prompt, self.strategy.assemble_response(chosen_tokens)
+        )
+        rejected_t, rejected_m = _encode_with_mask(
+            prompt, self.strategy.assemble_response(rejected_tokens)
+        )
+
+        return {
+            "chosen": chosen_t,
+            "chosen_mask": chosen_m,
+            "rejected": rejected_t,
+            "rejected_mask": rejected_m,
+        }
+
+    @property
+    def output_keys(self) -> List[str]:
+        return ["chosen", "chosen_mask", "rejected", "rejected_mask"]