include safety in app

app.py

@@ -5,19 +5,43 @@ import time
 import os
 import re
 from gradio_client import Client
+import torch
+from transformers import pipeline
 
-is_shared_ui = True if "fffiloni/consistent-character" in os.environ['SPACE_ID'] else False
-def safety_check(user_prompt, token):
+pipe_safety = pipeline("text-generation", model="HuggingFaceH4/zephyr-7b-beta", torch_dtype=torch.bfloat16, device_map="auto")
+
+agent_maker_sys = os.environ.get("SAFETY_PROMPT")
+
+instruction = f"""
+<|system|>
+{agent_maker_sys}</s>
+<|user|>
+"""
+
+def safety_check(user_prompt):
     
-    client = Client("fffiloni/safety-checker-bot", hf_token=token)
-    response = client.predict(
-            source_space="consistent-character space",
-            user_prompt=user_prompt,
-            api_name="/infer"
-    )
-    print(response)
+    prompt = f"{instruction.strip()}\n'{user_prompt}'</s>"   
+    print(f"""
+    —
+    USER PROMPT: {user_prompt}
+    """)
+    
+    outputs = pipe_safety(prompt, max_new_tokens=256, do_sample=True, temperature=0.3, top_k=50, top_p=0.95)
     
-    return response
+
+    pattern = r'\<\|system\|\>(.*?)\<\|assistant\|\>'
+    cleaned_text = re.sub(pattern, '', outputs[0]["generated_text"], flags=re.DOTALL)
+    
+    print(f"""
+    —
+    SAFETY COUNCIL: {cleaned_text}
+    
+    """)
+    
+    return cleaned_text.lstrip("\n")
+
+
+is_shared_ui = True if "fffiloni/consistent-character" in os.environ['SPACE_ID'] else False
     
 from utils.gradio_helpers import parse_outputs, process_outputs
 
@@ -33,9 +57,8 @@ def predict(request: gr.Request, *args, progress=gr.Progress(track_tqdm=True)):
     
     try:
         if is_shared_ui:
-            hf_token = os.environ.get("HF_TOKEN")
 
-            is_safe = safety_check(args[0], hf_token)
+            is_safe = safety_check(args[0])
             print(is_safe)
     
             match = re.search(r'\bYes\b', is_safe)