feat: enhance debiasing functionality with improved data handling and…

… bias detection logic

langtest/augmentation/__init__.py

@@ -1,4 +1,11 @@
 from .base import BaseAugmentaion, AugmentRobustness, TemplaticAugment
 from .augmenter import DataAugmenter
+from .debias import DebiasTextProcessing
 
-__all__ = ["BaseAugmentaion", "AugmentRobustness", "TemplaticAugment", "DataAugmenter"]
+__all__ = [
+    "DebiasTextProcessing",
+    "BaseAugmentaion",
+    "AugmentRobustness",
+    "TemplaticAugment",
+    "DataAugmenter",
+]

langtest/augmentation/debias.py

@@ -1,7 +1,9 @@
-from typing import Dict, List, Literal, Union
+from typing import Dict, List, Literal, TypeVar, Union
 from pydantic import BaseModel, Field
 import pandas as pd
 
+_Schema = TypeVar("_Schema")
+
 
 class BiasDetectionRequest(BaseModel):
     """
@@ -73,6 +75,10 @@ class DebiasingRequest(BaseModel):
     dataset: pd.DataFrame = Field(..., title="Dataset to debias")
     text_column: str = Field(..., title="Column name containing text")
 
+    model_config: Dict = {
+        "arbitrary_types_allowed": True,
+    }
+
 
 class DebiasingResult(BaseModel):
     """
@@ -86,52 +92,116 @@ class DebiasingResult(BaseModel):
     debiased_dataset: pd.DataFrame = Field(..., title="Debiased dataset")
     debias_info: List[Dict] = Field(..., title="Information about debiasing")
 
+    model_config: Dict = {
+        "arbitrary_types_allowed": True,
+    }
+
 
 class DebiasTextProcessing:
-    def __init__(self, dataset: pd.DataFrame, text_column: str):
-        self.dataset = dataset
-        self.text_column = text_column
-        self.debias_info = []
+    def __init__(self, model: str, hub: str, prompt: str):
+        # from langtest.tasks.task import TaskManager
+
+        # task = TaskManager("question-answering")
 
-    def initialize(self, model: str, hub: str):
-        # Placeholder for model initialization
-        self.debias_model = (model, hub)
+        # # self.debias_model = task.model(model_path=model, model_hub=hub, model_type="chat")
+        self.prompt = prompt
+        self.debias_info = pd.DataFrame(
+            {"row": [], "reason": [], "category": [], "sub_category": []}
+        )
+
+    def initialize(
+        self, input_dataset: pd.DataFrame, text_column: str, output_dataset: str = None
+    ):
+        self.input_dataset = input_dataset
+        self.text_column = text_column
+        self.output_dataset: pd.DataFrame = output_dataset
 
     def identify_bias(self):
-        for index, row in self.dataset.iterrows():
+        for index, row in self.input_dataset.iterrows():
             text = row[self.text_column]
-            reason, category, sub_category = self.detect_bias(text)
-            if reason:
-                self.debias_info.append(
-                    {
+            category, sub_category, rationale = self.detect_bias(text)
+            if rationale:
+                if index not in self.debias_info["row"].values:
+                    self.debias_info.loc[len(self.debias_info)] = {
                         "row": index,
-                        "reason": reason,
+                        "reason": rationale,
                         "category": category,
                         "sub_category": sub_category,
                     }
-                )
+                else:
+                    self.debias_info.loc[row["row"], "reason"] = rationale
+                    self.debias_info.loc[row["row"], "category"] = category
+                    self.debias_info.loc[row["row"], "sub_category"] = sub_category
+                self.debias_info = self.debias_info.reset_index(drop=True)
 
     def detect_bias(
         self, text: Union[str, BiasDetectionRequest]
     ) -> BiasDetectionResponse:
         # Placeholder for bias detection logic
+        if isinstance(text, BiasDetectionRequest):
+            text = text.text
+
+        output_data = self.interaction_llm(text, output_schema=BiasDetectionResponse)
+
+        return (
+            output_data.category,
+            output_data.sub_category,
+            output_data.bias_rationale,
+        )
 
-        return (None, None, None)
+    def interaction_llm(self, text: str, output_schema: type[_Schema]) -> _Schema:
+        import openai
 
-    def debias_text(self, text: str):
+        client = openai.Client()
+
+        response = client.beta.chat.completions.parse(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": self.prompt},
+                {"role": "user", "content": text},
+            ],
+            response_format=output_schema,
+        )
+
+        output_data = response.choices[0].message.parsed
+
+        return output_data
+
+    def debias_text(
+        self, text: str, category: str, sub_category: str, reason: str
+    ) -> str:
         # Placeholder for debiasing logic
-        return text
+        prompt = f"""
+        Debias the text with following bias information and reason.
+
+        Category: {category}
+        Sub-category: {sub_category}
+        Reason: {reason}
+
+        Text: {text}
+
+        Output:
+        """
+
+        debiased_text = self.interaction_llm(prompt, output_schema=DebiasedTextResponse)
+
+        return debiased_text.debiased_text
 
     def apply_debiasing(self):
-        for info in self.debias_info:
-            original_text = self.dataset.at[info["row"], self.text_column]
-            debiased_text = self.debias_text(original_text)
-            self.dataset.at[info["row"], self.text_column] = debiased_text
+        for idx, row in self.debias_info.iterrows():
+            original_text = self.input_dataset.at[row["row"], self.text_column]
+            debiased_text = self.debias_text(
+                original_text,
+                category=row["category"],
+                sub_category=row["sub_category"],
+                reason=row["reason"],
+            )
+            self.output_dataset.loc[row["row"], self.text_column] = debiased_text
 
     def process(self):
         self.identify_bias()
         self.apply_debiasing()
-        return self.dataset, self.debias_info
+        return self.output_dataset, self.debias_info
 
     def load_data(self, source: str, source_type: str):
         if source_type == "csv":