update

binoculars/detector.py

@@ -20,6 +20,9 @@ huggingface_config = {
 BINOCULARS_ACCURACY_THRESHOLD = 0.9015310749276843  # optimized for f1-score
 BINOCULARS_FPR_THRESHOLD = 0.8536432310785527  # optimized for low-fpr [chosen at 0.01%]
 
+DEVICE_1 = "cuda:0" if torch.cuda.is_available() else "cpu"
+DEVICE_2 = "cuda:1" if torch.cuda.device_count() > 1 else DEVICE_1
+
 
 class Binoculars(object):
     def __init__(self,
@@ -33,14 +36,14 @@ class Binoculars(object):
 
         self.change_mode(mode)
         self.observer_model = AutoModelForCausalLM.from_pretrained(observer_name_or_path,
-                                                                   device_map={"": "cuda" if torch.cuda.is_available() else "cpu"},
+                                                                   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={"": "cuda" if torch.cuda.is_available() else "cpu"},
+                                                                    device_map={"": DEVICE_2},
                                                                     trust_remote_code=True,
                                                                     torch_dtype=torch.bfloat16 if use_bfloat16
                                                                     else torch.float32,

model_utils.py

@@ -7,6 +7,8 @@ from NN_classifier.simple_binary_classifier import Medium_Binary_Network
 from feature_extraction import extract_features
 import pandas as pd
 
+DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
 def load_model(model_dir='models/medium_binary_classifier'):
     model_path = os.path.join(model_dir, 'nn_model.pt')
     scaler_path = os.path.join(model_dir, 'scaler.joblib')
@@ -27,8 +29,8 @@ def load_model(model_dir='models/medium_binary_classifier'):
     
     input_size = scaler.n_features_in_
     
-    model = Medium_Binary_Network(input_size, hidden_sizes=[256, 192, 128, 64], dropout=0.3).to("cuda" if torch.cuda.is_available() else "cpu")
-    model.load_state_dict(torch.load(model_path, map_location="cuda" if torch.cuda.is_available() else "cpu"))
+    model = Medium_Binary_Network(input_size, hidden_sizes=[256, 192, 128, 64], dropout=0.3).to(DEVICE)
+    model.load_state_dict(torch.load(model_path, map_location=DEVICE))
     model.eval()
     
     if imputer is not None:
@@ -74,7 +76,7 @@ def classify_text(text, model, scaler, label_encoder, imputer=None, scores=None)
     
     features_scaled = scaler.transform(features)
     
-    features_tensor = torch.FloatTensor(features_scaled).to("cuda" if torch.cuda.is_available() else "cpu")
+    features_tensor = torch.FloatTensor(features_scaled).to(DEVICE)
     
     with torch.no_grad():
         outputs = model(features_tensor)