implemented detection

wm_detector/core/detector.py

@@ -6,12 +6,9 @@
 
 import numpy as np
 from scipy import special
-from scipy.optimize import fminbound
 
 import torch
-from transformers import AutoTokenizer, LlamaForCausalLM
-
-import random
+from transformers import AutoTokenizer
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 

wm_detector/core/utils.py

@@ -1,6 +1,6 @@
 import random
 import numpy as np
-from .detector import AutoTokenizer, WmDetector
+from .detector import WmDetector
 
 def generate_pastel_color():
     """Generate a pastel color in HSL format."""

wm_detector/web/app.py

@@ -3,9 +3,22 @@ Main Flask application for the watermark detection web interface.
 """
 
 from flask import Flask, render_template, request, jsonify
-from ..core.detector import MarylandDetector, MarylandDetectorZ, OpenaiDetector, OpenaiDetectorZ, AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from ..core.detector import MarylandDetector, MarylandDetectorZ, OpenaiDetector, OpenaiDetectorZ
 from ..core.utils import get_token_details
 
+def convert_nan_to_null(obj):
+    """Convert NaN values to null for JSON serialization"""
+    import math
+    if isinstance(obj, float) and math.isnan(obj):
+        return None
+    elif isinstance(obj, dict):
+        return {k: convert_nan_to_null(v) for k, v in obj.items()}
+    elif isinstance(obj, list):
+        return [convert_nan_to_null(item) for item in obj]
+    return obj
+
 def create_detector(detector_type, tokenizer, **kwargs):
     """Create a detector instance based on the specified type."""
     detector_map = {
@@ -39,20 +52,11 @@ def create_app():
     # Add zip to Jinja's global context
     app.jinja_env.globals.update(zip=zip)
 
-    # Minimal setup: pick a detector (example: MarylandDetector)
-    model_id = "meta-llama/Llama-3.2-1B-Instruct"
+    # Pick a model
+    # model_id = "meta-llama/Llama-3.2-1B-Instruct"
+    model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id, cache_dir="wm_detector/static/hf_cache")
-
-    def convert_nan_to_null(obj):
-        """Convert NaN values to null for JSON serialization"""
-        import math
-        if isinstance(obj, float) and math.isnan(obj):
-            return None
-        elif isinstance(obj, dict):
-            return {k: convert_nan_to_null(v) for k, v in obj.items()}
-        elif isinstance(obj, list):
-            return [convert_nan_to_null(item) for item in obj]
-        return obj
+    model = AutoModelForCausalLM.from_pretrained(model_id)
 
     @app.route("/", methods=["GET"])
     def index():
@@ -68,9 +72,6 @@ def create_app():
             text = data.get('text', '')
             params = data.get('params', {})
             
-            if not isinstance(text, str):
-                return jsonify({'error': 'Text must be a string'}), 400
-            
             # Create a detector instance with the provided parameters
             detector = create_detector(
                 detector_type=params.get('detector_type', 'maryland'),