Update binoculars/detector.py

binoculars/detector.py

@@ -24,6 +24,7 @@ class Binoculars(object):
                  use_bfloat16: bool = True,
                  max_token_observed: int = 512,
                  mode: str = "low-fpr",
+                 quantize: bool = True
                  ) -> None:
         assert_tokenizer_consistency(observer_name_or_path, performer_name_or_path)
 
@@ -37,11 +38,20 @@ class Binoculars(object):
         # Load models with memory-efficient settings
         model_kwargs = {
             "device_map": "auto",
-            "load_in_8bit": True,
             "trust_remote_code": True,
             "token": huggingface_config["TOKEN"]
         }
 
+        if quantize:
+            try:
+                import bitsandbytes as bnb
+                model_kwargs["load_in_8bit"] = True
+            except ImportError:
+                print("bitsandbytes not available. Falling back to full precision.")
+                model_kwargs["torch_dtype"] = torch.bfloat16 if use_bfloat16 else torch.float32
+        else:
+            model_kwargs["torch_dtype"] = torch.bfloat16 if use_bfloat16 else torch.float32
+
         self.observer_model = AutoModelForCausalLM.from_pretrained(observer_name_or_path, **model_kwargs)
         self.performer_model = AutoModelForCausalLM.from_pretrained(performer_name_or_path, **model_kwargs)