Update app.py

app.py

@@ -7,182 +7,197 @@ import onnxruntime as ort
 from huggingface_hub import hf_hub_download
 import gradio as gr
 
+
 class ONNXInferencePipeline:
-    def __init__(self, repo_id):
-        # Retrieve the Hugging Face token from the environment variable
+    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")
-        if hf_token is None:
-            raise ValueError("HF_TOKEN environment variable is not set.")
 
         # Load banned keywords list
         self.banned_keywords = self.load_banned_keywords()
         print(f"Loaded {len(self.banned_keywords)} banned keywords")
-        
-        # Download files from Hugging Face Hub using the token
-        self.onnx_path = hf_hub_download(repo_id=repo_id, filename="model.onnx", use_auth_token=hf_token)
-        self.tokenizer_path = hf_hub_download(repo_id=repo_id, filename="train_bpe_tokenizer.json", use_auth_token=hf_token)
-        self.config_path = hf_hub_download(repo_id=repo_id, filename="hyperparameters.json", use_auth_token=hf_token)
+
+        # 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) as f:
+        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 = 256
+        self.max_len = int(self.config.get("max_len", 256))
 
         # Initialize ONNX runtime session
-        self.session = ort.InferenceSession(self.onnx_path)
-        self.providers = ['CPUExecutionProvider']  # Use CUDA if available
-        if 'CUDAExecutionProvider' in ort.get_available_providers():
-            self.providers = ['CUDAExecutionProvider']
-        self.session.set_providers(self.providers)
+        # 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):
-        # For testing purposes, using a small list
-        # In production, load your full list
-        
+        """
+        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:
-            raise Exception("Environment variable 'banned' is not set. Please set it with your banned keywords list.")
+            print("Environment variable 'banned' is not set. Using empty list.")
+            return []
 
-        # Create a local namespace to execute the code
+        # 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 = {}
-        
-        # Wrap the code in a function to allow return statements
         wrapped_code = f"""
 def get_banned_keywords():
 {textwrap.indent(code_str, '    ')}
 """
-        
         try:
-            # Execute the wrapped code
-            exec(wrapped_code, globals(), local_vars)
-            # Call the function to get the banned keywords
-            return local_vars['get_banned_keywords']()
+            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: {e}")
-            # Return a default empty list if there's an error
+            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."""
+        """Check if the input text contains any banned keywords as whole words."""
         text_lower = text.lower()
-        
-        # Split the text into words
-        words = ''.join(c if c.isalnum() else ' ' for c in text_lower).split()
-        
+        words = "".join(c if c.isalnum() else " " for c in text_lower).split()
+        word_set = set(words)
+
         for keyword in self.banned_keywords:
-            keyword_lower = keyword.lower()
-            
-            # Check if keyword is a whole word in the text
-            if keyword_lower in words:
+            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")
+        print("Keywords Passed. No inappropriate keywords found")
         return False
 
     def preprocess(self, text):
         encoding = self.tokenizer.encode(text)
-        ids = encoding.ids[:self.max_len]
+        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):
-        print(f"\nProcessing input: '{text[:50]}...' ({len(text)} characters)")
-        
-        # First check if the text contains any banned keywords
+        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': 'Inappropriate Content',
-                'probabilities': [1.0, 0.0]  # Assuming [inappropriate, appropriate]
+                "label": self.class_labels[0],
+                "probabilities": [1.0, 0.0] if len(self.class_labels) == 2 else [1.0] * len(self.class_labels),
             }
-        
-        # If no banned keywords found, proceed with model prediction
-        print("Running ML model for classification...")
-        
+
         # Preprocess
         input_array = self.preprocess(text)
 
-        # Run inference
-        results = self.session.run(
-            None,
-            {'input': input_array}
-        )
+        # Run inference. Use detected input name
+        outputs = self.session.run(None, {self.input_name: input_array})
 
-        # Post-process
-        logits = results[0]
-        probabilities = np.exp(logits) / np.sum(np.exp(logits), axis=1, keepdims=True)
-        predicted_class = int(np.argmax(probabilities))
-
-        # Map to labels
-        class_labels = ['Inappropriate Content', 'Appropriate']
-        
-        # Log model result
-        print(f"Model Passed - Result: {class_labels[predicted_class]} (Confidence: {probabilities[0][predicted_class]:.2%})")
-        
-        return {
-            'label': class_labels[predicted_class],
-            'probabilities': probabilities[0].tolist()
-        }
-
-# Example usage
-if __name__ == "__main__":
-    # Initialize the pipeline with the Hugging Face repository ID
-    print("Initializing content filter pipeline...")
-    pipeline = ONNXInferencePipeline(repo_id="iimran/abuse-detector")
-    print("Pipeline initialized successfully")
-
-    # Example texts for testing
-    example_texts = [
-        "You're a worthless piece of garbage who should die",
-        "Hello HR, I hope this message finds you well. I'm 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."
-    ]
-    
-    for text in example_texts:
-        result = pipeline.predict(text)
-        print(f"Input: {text[:50]}...")
-        print(f"Prediction: {result['label']} ")
-        print("-" * 80)
-
-    # Define a function for Gradio to use
-    def gradio_predict(text):
-        result = pipeline.predict(text)
-        return (
-            f"Prediction: {result['label']} \n"
-        )
+        # 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()}
 
-    # Create a Gradio interface
+
+# 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 is designed to identify inappropriate content in text. "
-        "It analyzes input for Australian Slang language and abuses. "
-        "While it's trained on a compact dataset and may not catch highly nuanced or sophisticated language, "
-        "it effectively detects day-to-day offensive language commonly used in conversations."
-    ),
+            "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've outdone yourself in stupidity. A real cock-up of a human—should we clap for your bollocks-faced greatness or just pity you?",
-    "You're a mad bastard, but I'd still grab a beer with you! Fuck 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'm 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 don't hesitate to reach out if you need additional information—I'd be happy to discuss further. Looking forward to hearing from you soon!",
-    "The weather today is lovely, and I'm looking forward to a productive day at work.",
-
-    # Mixed examples (some offensive, some neutral)
-    "I appreciate your help, but honestly, you're such a clueless idiot sometimes. Still, thanks for trying."
-]
+            # 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."
+        ],
     )
-
-    # Launch the Gradio app
     print("Launching Gradio interface...")
-    # at the very bottom where you launch Gradio
     iface.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))