Merge pull request #2 from guardrails-ai/aaravnavani-patch-1

Update __init__.py

Author: dtam

inference/download_model.py

@@ -38,7 +38,7 @@
 )
 def download_model(model_name, model_revision, force_download=False):
     from huggingface_hub import snapshot_download
-    from transformers import AutoTokenizer, AutoModelForCausalLM
+    from transformers import AutoTokenizer
     import json
     import os
 

inference/serving-non-optimized.py

@@ -56,8 +56,8 @@ def generate(self, chat, is_user_prompt=True, enforce_policies=None):
         tokenizer = self.tokenizer
         model = self.model
 
-        print(f"Model: Loaded on device")
-        print(f"Model: Chat {chat}")
+        print("Model: Loaded on device")
+        print("Model: Chat {chat}")
 
         INPUT_POLICIES = {
             "NO_DANGEROUS_CONTENT": "\"No Dangerous Content\": The prompt shall not contain or seek generation of content that harming oneself and/or others (e.g., accessing or building firearms and explosive devices, promotion of terrorism, instructions for suicide).",

pyproject.toml

@@ -27,4 +27,9 @@ testpaths = [
 ]
 
 [tool.pyright]
-include = ["validator"]
+include = ["validator"]
+
+[tool.setuptools]
+packages = [
+    "validator"
+]

tests/test_validator.py

@@ -3,13 +3,13 @@
 
 from guardrails import Guard
 import pytest
-from validator import ValidatorTemplate
+from validator import ShieldGemma2B
 
 # We use 'exception' as the validator's fail action,
 #  so we expect failures to always raise an Exception
 # Learn more about corrective actions here:
 #  https://www.guardrailsai.com/docs/concepts/output/#%EF%B8%8F-specifying-corrective-actions
-guard = Guard.from_string(validators=[ValidatorTemplate(arg_1="arg_1", arg_2="arg_2", on_fail="exception")])
+guard = Guard.from_string(validators=[ShieldGemma2B(arg_1="arg_1", arg_2="arg_2", on_fail="exception")])
 
 def test_pass():
   test_output = "pass"

validator/__init__.py

@@ -1,3 +1,3 @@
-from .main import ValidatorTemplate
+from .main import ShieldGemma2B
 
-__all__ = ["ValidatorTemplate"]
+__all__ = ["ShieldGemma2B"]