Adding implementation for edit_similarity_node. (#94)

Summary: Pull Request resolved: #94

Differential Revision: D87895834

Author: anilkram

privacy_guard/analysis/extraction/edit_similarity_node.py

@@ -12,14 +12,19 @@
 # pyre-strict
 
 from dataclasses import dataclass
-from typing import Optional
+from typing import cast, Optional
 
 import pandas as pd
+import textdistance
 from privacy_guard.analysis.base_analysis_node import BaseAnalysisNode
 from privacy_guard.analysis.base_analysis_output import BaseAnalysisOutput
 from privacy_guard.analysis.extraction.text_inclusion_analysis_input import (
     TextInclusionAnalysisInput,
 )
+from privacy_guard.analysis.extraction.text_inclusion_analysis_node import (
+    _clean_text,
+    _normalize_by_target_len,
+)
 
 from tqdm import tqdm
 
@@ -59,12 +64,47 @@ def __init__(self, analysis_input: TextInclusionAnalysisInput) -> None:
         self.prompt_key: str = analysis_input.prompt_key
         self.generation_key: str = analysis_input.generation_key
         self.target_key: str = analysis_input.target_key
+        self.generation_df: pd.DataFrame = analysis_input.generation_df
         super().__init__(analysis_input=analysis_input)
 
+    def _compute_edit_similarity(
+        self, row: pd.Series, s1_column: str | None = None, s2_column: str | None = None
+    ) -> int:
+        """Compute edit similarity between target and generation text. Texts are cleaned first.
+        Currently not supported for multi target mode.
+
+        Args:
+            row (pd.Series): A row of a DataFrame containing the s1 and s2 columns.
+
+        Returns:
+            int: Edit similarity between the two strings.
+        """
+        s1 = _clean_text(row[s1_column or self.target_key])
+        s2 = _clean_text(row[s2_column or self.generation_key])
+        levenshtein = textdistance.levenshtein.similarity(s1, s2)
+        return levenshtein
+
     def run_analysis(self) -> EditSimilarityNodeOutput:
-        return EditSimilarityNodeOutput(
-            num_samples=0,
+        analysis_input: TextInclusionAnalysisInput = cast(
+            TextInclusionAnalysisInput, self.analysis_input
+        )
+        generation_df = analysis_input.generation_df
+
+        outputs = EditSimilarityNodeOutput(
+            num_samples=len(generation_df),
             edit_similarity=None,
             edit_similarity_score=None,
-            augmented_output_dataset=pd.DataFrame(),
+            augmented_output_dataset=generation_df,
+        )
+
+        generation_df["edit_similarity"] = generation_df.progress_apply(
+            self._compute_edit_similarity, axis=1
         )
+        generation_df["edit_similarity_score"] = _normalize_by_target_len(
+            generation_df["edit_similarity"], generation_df["target"]
+        )
+
+        outputs.edit_similarity = generation_df["edit_similarity"]
+        outputs.edit_similarity_score = generation_df["edit_similarity_score"]
+
+        return outputs

privacy_guard/analysis/tests/test_edit_similarity_node.py

@@ -74,5 +74,11 @@ def test_edit_similarity_node_output_creation(self) -> None:
         self.assertEqual(output.augmented_output_dataset.equals(generation_df), True)
 
     def test_text_inclusion_edit_similarity(self) -> None:
-        analysis_node = EditSimilarityNode(analysis_input=self.analysis_input)
-        analysis_node.run_analysis()
+        analysis_input = TextInclusionAnalysisInput(
+            generation_df=pd.DataFrame(self.data)
+        )
+        analysis_node = EditSimilarityNode(analysis_input=analysis_input)
+        results = analysis_node.compute_outputs()
+        self.assertIn("edit_similarity", results)
+        self.assertIn("edit_similarity_score", results)
+        self.assertEqual(results["edit_similarity"].tolist(), [13, 3, 22, 16, 8, 16])