Add self similarity classifier

aisploit/classifiers/__init__.py

@@ -1,11 +1,13 @@
 from .markdown import MarkdownInjectionClassifier
 from .package_hallucination import PythonPackageHallucinationClassifier
+from .self_similarity import SelfSimilarityClassifier
 from .text import RegexClassifier, SubstringClassifier, TextTokenClassifier
 
 __all__ = [
     "MarkdownInjectionClassifier",
     "PythonPackageHallucinationClassifier",
     "RegexClassifier",
     "SubstringClassifier",
+    "SelfSimilarityClassifier",
     "TextTokenClassifier",
 ]

aisploit/classifiers/self_similarity.py

@@ -0,0 +1,50 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.util import cos_sim
+
+from ..core import BaseTextClassifier, Score
+
+
+@dataclass
+class SelfSimilarityClassifier(BaseTextClassifier[float]):
+    """A text classifier based on self-similarity using cosine similarity scores."""
+
+    model_name_or_path: str = "all-MiniLM-L6-v2"
+    threshold: float = 0.7
+    tags: List[str] = field(default_factory=lambda: ["hallucination"], init=False)
+
+    def __post_init__(self) -> None:
+        """Initialize the SentenceTransformer model."""
+        self._model = SentenceTransformer(self.model_name_or_path)
+
+    def score(self, input: str, references: List[str] | None = None) -> Score[float]:
+        """Score the input text based on its self-similarity to reference texts.
+
+        Args:
+            input (str): The input text to be scored.
+            references (List[str], optional): List of reference texts. Defaults to None.
+
+        Raises:
+            ValueError: If references is None or if the number of references is not at least 1.
+
+        Returns:
+            Score[float]: A Score object representing the self-similarity score of the input.
+        """
+        if not references or not len(references) >= 1:
+            raise ValueError("The number of references must be at least 1.")
+
+        input_embeddings = self._model.encode(input, convert_to_tensor=True)
+        references_embeddings = self._model.encode(references, convert_to_tensor=True)
+
+        cos_scores = cos_sim(input_embeddings, references_embeddings)[0]
+
+        score = cos_scores.mean()
+
+        return Score[float](
+            flagged=(score < self.threshold).item(),
+            value=score.item(),
+            description="Returns True if the cosine similarity score is less than the threshold",
+            explanation=f"The cosine similarity score for the input is {score}",
+        )

examples/classifier.ipynb

@@ -26,6 +26,7 @@
    "source": [
     "import textwrap\n",
     "from dotenv import load_dotenv\n",
+    "from aisploit.classifiers import SelfSimilarityClassifier\n",
     "from aisploit.classifiers.presidio import PresidioAnalyserClassifier\n",
     "from aisploit.classifiers.huggingface import (\n",
     "    BleuClassifier,\n",
@@ -41,6 +42,34 @@
     "load_dotenv()"
    ]
   },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "Score(flagged=True, value=0.20951247215270996, description='Returns True if the cosine similarity score is less than the threshold', explanation='The cosine similarity score for the input is 0.20951247215270996')"
+      ]
+     },
+     "execution_count": 2,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "classifier = SelfSimilarityClassifier()\n",
+    "classifier.score(\n",
+    "    \"The sky turned green and the trees began to sing in a chorus of laughter.\", \n",
+    "    [\n",
+    "        \"As I looked around, I noticed that the buildings were made of candy, and the streets were paved with shimmering gold.\",\n",
+    "        \"I found myself in a library unlike any other, where the books flew off the shelves and started telling stories of their own.\",\n",
+    "        \"Every person I met had wings, and they soared through the air with grace and elegance, leaving trails of glitter in their wake.\",\n",
+    "    ],\n",
+    ")"
+   ]
+  },
   {
    "cell_type": "markdown",
    "metadata": {},

pyproject.toml

@@ -53,6 +53,7 @@ pillow = "^10.3.0"
 tqdm = "^4.66.2"
 evaluate = "^0.4.1"
 bert-score = "^0.3.13"
+sentence-transformers = "^2.7.0"
 
 [tool.poetry.group.dev.dependencies]
 chromadb = "^0.4.23"