Upload 3 files

model_inference/gaurdrails_manager.py

@@ -0,0 +1,64 @@
+from endpoints.api_models import OutputGuardrailsConfig , LLMResponse
+from model_inference.groundedness_checker import GroundednessChecker
+import re
+
+# 
+groundedness_checker = GroundednessChecker(model_path="./grounding_detector")
+
+# A simple result class to hold individual check outcomes.
+class Result:
+    def __init__(self):
+        self.details = {}
+
+    def add(self, rule_name: str, passed: bool):
+        self.details[rule_name] = passed
+
+    def grounded(self) -> bool:
+        # The response is considered "grounded" if all enabled rules pass.
+        return all(self.details.values())
+
+class ContextualGroundednessCheck:
+    name = "Contextual Groundedness"
+    
+    def check(self,llm_response:LLMResponse) -> bool:
+        groundedness_check = groundedness_checker.check(llm_response.question, llm_response.answer, llm_response.context)
+        print(groundedness_check)
+        return groundedness_check['is_grounded']
+
+
+class ToxicityRule:
+
+    name = "Toxicity"
+
+    def check(self, llm_response:LLMResponse) -> bool:
+        
+        no_toxicity = True
+        matched =  re.search(r"(hate|kill|suicide|selfharm)", llm_response.answer, re.IGNORECASE)
+        
+        if matched:
+            no_toxicity = False
+        
+        return no_toxicity
+    
+
+# Manager class to load and execute the enabled guardrail rules.
+class GuardrailsManager:
+    def __init__(self, config: OutputGuardrailsConfig):
+        self.config = config
+        self.rules = self.load_rules()
+
+    def load_rules(self):
+        rules = []
+        if self.config.contextual_grounding:
+            rules.append(ContextualGroundednessCheck())
+        if self.config.toxicity:
+            rules.append(ToxicityRule())
+        # Add additional rules based on configuration here.
+        return rules
+
+    def check(self, llm_response: LLMResponse) -> Result:
+        result = Result()
+        for rule in self.rules:
+            rule_result = rule.check(llm_response)
+            result.add(rule.name, rule_result)
+        return result

model_inference/groundedness_checker.py

@@ -0,0 +1,64 @@
+from transformers import AutoTokenizer, AutoModelForSequenceClassification
+import torch
+
+class GroundednessChecker:
+    def __init__(self, model_path="./grounding_detector"):
+        self.tokenizer = AutoTokenizer.from_pretrained(model_path)
+        self.model = AutoModelForSequenceClassification.from_pretrained(model_path)
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model.to(self.device)
+        
+    def check(self, question: str, answer: str, context: str) -> dict:
+        """Check if answer is grounded in context"""
+        inputs = self.tokenizer(
+            question,
+            answer + " [SEP] " + context,
+            padding=True,
+            truncation=True,
+            max_length=512,
+            return_tensors="pt"
+        ).to(self.device)
+        
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        
+        probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
+        return {
+            "is_grounded": bool(torch.argmax(probs)),
+            "confidence": probs[0][1].item(),
+            "details": {
+                "question": question,
+                "answer": answer,
+                "context_snippet": context[:200] + "..." if len(context) > 200 else context
+            }
+        }
+
+# Usage Example
+if __name__ == "__main__":
+    # Initialize checker
+    checker = GroundednessChecker()
+    
+    # Example from banking PDS
+    context = """
+    Premium Savings Account Terms:
+    - Annual Percentage Yield (APY): 4.25%
+    - Minimum opening deposit: $1,000
+    - Monthly maintenance fee: $5 (waived if daily balance >= $1,000)
+    - Maximum withdrawals: 6 per month
+    """
+    
+    # Grounded example
+    grounded_result = checker.check(
+        question="What is the minimum opening deposit?",
+        answer="$1,000",
+        context=context
+    )
+    print("Grounded Result:", grounded_result)
+
+    # Ungrounded example
+    ungrounded_result = checker.check(
+        question="What is the monthly maintenance fee?",
+        answer="$10 monthly charge",
+        context=context
+    )
+    print("Ungrounded Result:", ungrounded_result)