Update binoculars/detector.py

binoculars/detector.py

@@ -15,9 +15,7 @@ torch.set_grad_enabled(False)
 BINOCULARS_ACCURACY_THRESHOLD = 0.9015310749276843  # optimized for f1-score
 BINOCULARS_FPR_THRESHOLD = 0.8536432310785527  # optimized for low-fpr
 
-DEVICE_1 = "cuda:0" if torch.cuda.is_available() else "cpu"
-DEVICE_2 = "cuda:1" if torch.cuda.device_count() > 1 else DEVICE_1
-
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
 class Binoculars(object):
     def __init__(self,
@@ -36,20 +34,16 @@ class Binoculars(object):
         else:
             raise ValueError(f"Invalid mode: {mode}")
 
-        self.observer_model = AutoModelForCausalLM.from_pretrained(observer_name_or_path,
-                                                                   device_map={"": DEVICE_1},
-                                                                   trust_remote_code=True,
-                                                                   torch_dtype=torch.bfloat16 if use_bfloat16
-                                                                   else torch.float32,
-                                                                   token=huggingface_config["TOKEN"]
-                                                                   )
-        self.performer_model = AutoModelForCausalLM.from_pretrained(performer_name_or_path,
-                                                                    device_map={"": DEVICE_2},
-                                                                    trust_remote_code=True,
-                                                                    torch_dtype=torch.bfloat16 if use_bfloat16
-                                                                    else torch.float32,
-                                                                    token=huggingface_config["TOKEN"]
-                                                                    )
+        # Load models with memory-efficient settings
+        model_kwargs = {
+            "device_map": "auto",
+            "load_in_8bit": True,
+            "trust_remote_code": True,
+            "token": huggingface_config["TOKEN"]
+        }
+
+        self.observer_model = AutoModelForCausalLM.from_pretrained(observer_name_or_path, **model_kwargs)
+        self.performer_model = AutoModelForCausalLM.from_pretrained(performer_name_or_path, **model_kwargs)
 
         self.observer_model.eval()
         self.performer_model.eval()
@@ -76,15 +70,13 @@ class Binoculars(object):
             padding="longest" if batch_size > 1 else False,
             truncation=True,
             max_length=self.max_token_observed,
-            return_token_type_ids=False).to(self.observer_model.device)
+            return_token_type_ids=False)
         return encodings
 
     @torch.inference_mode()
     def _get_logits(self, encodings: transformers.BatchEncoding) -> torch.Tensor:
-        observer_logits = self.observer_model(**encodings.to(DEVICE_1)).logits
-        performer_logits = self.performer_model(**encodings.to(DEVICE_2)).logits
-        if DEVICE_1 != "cpu":
-            torch.cuda.synchronize()
+        observer_logits = self.observer_model(**encodings).logits
+        performer_logits = self.performer_model(**encodings).logits
         return observer_logits, performer_logits
 
     def compute_score(self, input_text: Union[list[str], str]) -> Union[float, list[float]]:
@@ -92,8 +84,7 @@ class Binoculars(object):
         encodings = self._tokenize(batch)
         observer_logits, performer_logits = self._get_logits(encodings)
         ppl = perplexity(encodings, performer_logits)
-        x_ppl = entropy(observer_logits.to(DEVICE_1), performer_logits.to(DEVICE_1),
-                        encodings.to(DEVICE_1), self.tokenizer.pad_token_id)
+        x_ppl = entropy(observer_logits, performer_logits, encodings, self.tokenizer.pad_token_id)
         binoculars_scores = ppl / x_ppl
         binoculars_scores = binoculars_scores.tolist()
         return binoculars_scores[0] if isinstance(input_text, str) else binoculars_scores
@@ -104,4 +95,4 @@ class Binoculars(object):
                         "Most likely AI-generated",
                         "Most likely human-generated"
                         ).tolist()
-        return pred
+        return pred[0] if isinstance(input_text, str) else pred