Setup ML Endpoints

app.py

@@ -1,7 +1,5 @@
-import json
-import torch
 import nltk
-from typing import Any, Dict, List
+from typing import Any, Dict
 import spacy
 class InferlessPythonModel:
 

validator/main.py

@@ -1,10 +1,11 @@
-import nltk
-import spacy
+import json
 import re
-from typing import Callable, Dict, List, Optional
+from typing import Any, Callable, Dict, List, Optional
 
-# from guardrails.validator_base import ErrorSpan
+import nltk
+import spacy
 from guardrails.logger import logger
+from guardrails.validator_base import ErrorSpan
 from guardrails.validators import (
     FailResult,
     PassResult,
@@ -14,7 +15,9 @@
 )
 
 
-@register_validator(name="guardrails/competitor_check", data_type="string")
+@register_validator(
+    name="guardrails/competitor_check", data_type="string", has_guardrails_endpoint=True
+)
 class CompetitorCheck(Validator):
     """Validates that LLM-generated text is not naming any competitors from a
     given list.
@@ -56,8 +59,13 @@ def __init__(
         self,
         competitors: List[str],
         on_fail: Optional[Callable] = None,
+        **kwargs,
     ):
-        super().__init__(competitors=competitors, on_fail=on_fail)
+        super().__init__(
+            competitors=competitors,
+            on_fail=on_fail,
+            **kwargs,
+        )
         self._competitors = competitors
         model = "en_core_web_trf"
         self.nlp = spacy.load(model)
@@ -82,19 +90,18 @@ def exact_match(self, text: str, competitors: List[str]) -> List[str]:
                 found_entities.append(entity)
         return found_entities
 
-    def perform_ner(self, text: str, nlp) -> List[str]:
+    def perform_ner(self, text: str) -> List[str]:
         """Performs named entity recognition on text using a provided NLP
         model.
 
         Args:
             text (str): The text to perform named entity recognition on.
-            nlp: The NLP model to use for entity recognition.
 
         Returns:
             entities: A list of entities found.
         """
 
-        doc = nlp(text)
+        doc = self.nlp(text)
         entities = []
         for ent in doc.ents:
             entities.append(ent.text)
@@ -138,14 +145,17 @@ def validate(self, value: str, metadata=Dict) -> ValidationResult:
         sentences = nltk.sent_tokenize(value)
         flagged_sentences = []
         filtered_sentences = []
-        # error_spans:List[ErrorSpan] = []
+        error_spans: List[ErrorSpan] = []
         list_of_competitors_found = []
         start_ind = 0
         for sentence in sentences:
-
             entities = self.exact_match(sentence, self._competitors)
             if entities:
-                ner_entities = self.perform_ner(sentence, self.nlp)
+                ner_entities = self._inference(
+                    {"text": sentence, "competitors": self._competitors}
+                )
+                if isinstance(ner_entities, str):
+                    ner_entities = [ner_entities]
                 found_competitors = self.is_entity_in_list(ner_entities, entities)
                 if found_competitors:
                     flagged_sentences.append((found_competitors, sentence))
@@ -164,20 +174,40 @@ def validate(self, value: str, metadata=Dict) -> ValidationResult:
             for entity in tup[0]:
                 found_entities.append(entity)
 
-        # def find_all(a_str, sub):
-        #     start = 0
-        #     while True:
-        #         start = a_str.find(sub, start)
-        #         if start == -1: 
-        #             return
-        #         yield start
-        #         start += len(sub) # use start += 1 to find overlapping matches
-
-        # error_spans = []
-        # for entity in found_entities: 
-        #     starts = list(find_all(value, entity))
-        #     for start in starts:
-        #         error_spans.append(ErrorSpan(start=start, end=start+len(entity), reason=f'Competitor found: {value[start:start+len(entity)]}'))
+        def find_all(a_str, sub):
+            start = 0
+            while True:
+                start = a_str.find(sub, start)
+                if start == -1:
+                    return
+                yield start
+                start += len(sub)  # use start += 1 to find overlapping matches
+
+        error_spans = []
+        for entity in found_entities:
+            starts = list(find_all(value, entity))
+            for start in starts:
+                error_spans.append(
+                    ErrorSpan(
+                        start=start,
+                        end=start + len(entity),
+                        reason=f"Competitor found: {value[start:start+len(entity)]}",
+                    )
+                )
+        def find_all(a_str, sub):
+            start = 0
+            while True:
+                start = a_str.find(sub, start)
+                if start == -1: 
+                    return
+                yield start
+                start += len(sub) # use start += 1 to find overlapping matches
+
+        error_spans = []
+        for entity in found_entities: 
+            starts = list(find_all(value, entity))
+            for start in starts:
+                error_spans.append(ErrorSpan(start=start, end=start+len(entity), reason=f'Competitor found: {value[start:start+len(entity)]}'))
 
 
         if len(flagged_sentences):
@@ -187,7 +217,46 @@ def validate(self, value: str, metadata=Dict) -> ValidationResult:
                     "Please avoid naming those competitors next time"
                 ),
                 fix_value=filtered_output,
-                # error_spans=error_spans
+                error_spans=error_spans,
             )
         else:
-            return PassResult()
+            return PassResult()
+
+    def _inference_local(self, model_input: Any) -> str:
+        """Local inference method to detect and anonymize competitor names."""
+        text = model_input["text"]
+        competitors = model_input["competitors"]
+
+        doc = self.nlp(text)
+        anonymized_text = text
+        for ent in doc.ents:
+            if ent.text in competitors:
+                anonymized_text = anonymized_text.replace(ent.text, "[COMPETITOR]")
+        return anonymized_text
+
+    def _inference_remote(self, model_input: Any) -> str:
+        """Remote inference method for a hosted ML endpoint."""
+        request_body = {
+            "inputs": [
+                {
+                    "name": "text",
+                    "shape": [1],
+                    "data": [model_input["text"]],
+                    "datatype": "BYTES"
+                },
+                {
+                    "name": "competitors",
+                    "shape": [len(model_input["competitors"])],
+                    "data": model_input["competitors"],
+                    "datatype": "BYTES"
+                }
+            ]
+        }
+        response = self._hub_inference_request(json.dumps(request_body), self.validation_endpoint)
+
+        if not response or "outputs" not in response:
+            raise ValueError("Invalid response from remote inference", response)
+
+        outputs = response["outputs"][0]["data"][0]
+
+        return outputs