import os import json import numpy as np import textwrap from tokenizers import Tokenizer import onnxruntime as ort from huggingface_hub import hf_hub_download import gradio as gr class ONNXInferencePipeline: def __init__(self, repo_id, repo_type="model"): # Read token from env. In a Space, HF_TOKEN can be set in the Secrets panel. hf_token = os.getenv("HF_TOKEN") # Load banned keywords list self.banned_keywords = self.load_banned_keywords() print(f"Loaded {len(self.banned_keywords)} banned keywords") # Download artifacts. Newer huggingface_hub uses token=, not use_auth_token= self.onnx_path = hf_hub_download( repo_id=repo_id, filename="model.onnx", token=hf_token, repo_type=repo_type ) self.tokenizer_path = hf_hub_download( repo_id=repo_id, filename="train_bpe_tokenizer.json", token=hf_token, repo_type=repo_type ) self.config_path = hf_hub_download( repo_id=repo_id, filename="hyperparameters.json", token=hf_token, repo_type=repo_type ) # Load configuration with open(self.config_path, "r") as f: self.config = json.load(f) # Initialize tokenizer self.tokenizer = Tokenizer.from_file(self.tokenizer_path) self.max_len = int(self.config.get("max_len", 256)) # Initialize ONNX runtime session # Spaces CPU runtime typically uses CPUExecutionProvider providers = ort.get_available_providers() if "CUDAExecutionProvider" in providers: use_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] else: use_providers = ["CPUExecutionProvider"] sess_options = ort.SessionOptions() # Reduce memory and improve cold start a bit sess_options.enable_mem_pattern = False sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL self.session = ort.InferenceSession(self.onnx_path, sess_options=sess_options, providers=use_providers) # Cache model input name to avoid mismatches like input vs input_ids self.input_name = self.session.get_inputs()[0].name print(f"ONNX model input name detected: {self.input_name}") # If you want label order from config, you can read it self.class_labels = self.config.get("class_labels", ["Inappropriate Content", "Appropriate"]) def load_banned_keywords(self): """ Load banned keywords from env var named 'banned'. Supports two formats: 1) Python code snippet that returns a list (your current method) 2) JSON array of strings """ code_str = os.getenv("banned") if not code_str: print("Environment variable 'banned' is not set. Using empty list.") return [] # Try JSON first try: parsed = json.loads(code_str) if isinstance(parsed, list) and all(isinstance(x, str) for x in parsed): return parsed except Exception: pass # Fallback to executable code that returns a list local_vars = {} wrapped_code = f""" def get_banned_keywords(): {textwrap.indent(code_str, ' ')} """ try: exec(wrapped_code, {}, local_vars) result = local_vars["get_banned_keywords"]() if isinstance(result, list): return [str(x) for x in result] print("Loaded banned keywords code did not return a list. Using empty list.") return [] except Exception as e: print(f"Error loading banned keywords from code: {e}") return [] def contains_banned_keyword(self, text): """Check if the input text contains any banned keywords as whole words.""" text_lower = text.lower() words = "".join(c if c.isalnum() else " " for c in text_lower).split() word_set = set(words) for keyword in self.banned_keywords: kw = str(keyword).lower().strip() if not kw: continue if kw in word_set: print(f"Keyword detected: '{keyword}'") return True print("Keywords Passed. No inappropriate keywords found") return False def preprocess(self, text): encoding = self.tokenizer.encode(text) ids = encoding.ids[: self.max_len] padding = [0] * (self.max_len - len(ids)) return np.array(ids + padding, dtype=np.int64).reshape(1, -1) @staticmethod def softmax(logits): # Numerically stable softmax x = logits - np.max(logits, axis=1, keepdims=True) e = np.exp(x) return e / np.sum(e, axis=1, keepdims=True) def predict(self, text): snippet = text[:50].replace("\n", " ") print(f"\nProcessing input: '{snippet}...' ({len(text)} characters)") # First rule based filter if self.contains_banned_keyword(text): print("Input rejected by keyword filter") return { "label": self.class_labels[0], "probabilities": [1.0, 0.0] if len(self.class_labels) == 2 else [1.0] * len(self.class_labels), } # Preprocess input_array = self.preprocess(text) # Run inference. Use detected input name outputs = self.session.run(None, {self.input_name: input_array}) # Post process logits = outputs[0] probs = self.softmax(logits) pred_idx = int(np.argmax(probs)) label = self.class_labels[pred_idx] if pred_idx < len(self.class_labels) else str(pred_idx) print(f"Model Passed. Result: {label} (Confidence: {probs[0][pred_idx]:.2%})") return {"label": label, "probabilities": probs[0].tolist()} # Gradio glue def gradio_predict(text): result = PIPELINE.predict(text) return f"Prediction: {result['label']}\n" # Create pipeline at import so the Space is ready print("Initializing content filter pipeline...") PIPELINE = ONNXInferencePipeline(repo_id="iimran/abuse-detector", repo_type="model") print("Pipeline initialized successfully") if __name__ == "__main__": # Required in Spaces. PORT is injected. Bind to 0.0.0.0 iface = gr.Interface( fn=gradio_predict, inputs=gr.Textbox(lines=7, placeholder="Enter text here..."), outputs="text", title="Abuse Detector - Offensive Language Detector", description=( "Abuse detector identifies inappropriate content in text. " "It analyzes input for Australian slang and abusive language. " "It is trained on a compact dataset. It may not catch highly nuanced language, " "but it detects common day to day offensive language." ), examples=[ # Explicitly offensive examples "Congrats, you fuckbrain arsehole, you have outdone yourself in stupidity. A real cock up of a human. Should we clap for your bollocks faced greatness or just pity you?", "You are a mad bastard, but I would still grab a beer with you. Mess around all you like, you cockheaded legend. Your arsehole antics are bloody brilliant.", "Your mother should have done better raising such a useless idiot.", # Neutral or appropriate examples "Hello HR, I hope this message finds you well. I am writing to express my gratitude for the opportunity to interview for the Financial Analyst position last week. It was a pleasure to meet you and learn more about the role and your team.", "Thank you for your time and consideration. Please reach out if you need anything. I would be happy to discuss further.", "The weather today is lovely, and I am looking forward to a productive day at work.", # Mixed "I appreciate your help, but honestly, you are such a clueless idiot sometimes. Still, thanks for trying." ], ) print("Launching Gradio interface...") iface.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))