DivEye - PR (fix bugs, cuda) (#8)

- cuda fixes (56e3ef5b1c6336e84d881fe737fd841db741d891)
- cuda fixes (b66e8e7aa1bc538ca08e64c1f45dd842a71f1762)

software.py

@@ -98,26 +98,19 @@ class BiScope:
 
 class Software:
     def __init__(self):
-        self.device_div = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        if torch.cuda.device_count() > 1:
-            self.device_bi = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
-        else:
-            self.device_bi = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-
         self.token = os.getenv("HF_TOKEN") 
 
         self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
         self.div_model = AutoModelForCausalLM.from_pretrained(
-            "tiiuae/falcon-7b", device_map=self.device_div, torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
+            "tiiuae/falcon-7b", torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
         )
 
         self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True, use_auth_token=self.token)
         self.bi_model = AutoModelForCausalLM.from_pretrained(
-            "google/gemma-1.1-2b-it", device_map=self.device_bi, torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
+            "google/gemma-1.1-2b-it", torch_dtype=torch.float16, trust_remote_code=True, use_auth_token=self.token
         )
         
-        self.diveye = Diversity(self.div_model, self.div_tokenizer, self.device_div)
-        self.biscope = BiScope(self.bi_model, self.bi_tokenizer, self.device_bi)
+        
         self.model_path = Path(__file__).parent / "model.json"
 
         self.model = xgb.XGBClassifier()
@@ -134,8 +127,21 @@ class Software:
     
     @spaces.GPU
     def evaluate(self, text):
-        diveye_features = self.diveye.compute_features(text)
-        biscope_features = self.biscope.detect_single_sample(text)
+        # Load models to GPUs.
+        device_div = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        if torch.cuda.device_count() > 1:
+            device_bi = torch.device("cuda:1" if torch.cuda.is_available() else "cpu")
+
+        if not next(self.div_model.parameters()).is_cuda:
+            self.div_model = self.div_model.to(device_div)
+        if not next(self.bi_model.parameters()).is_cuda:
+            self.bi_model = self.bi_model.to(device_bi)
+            
+        diveye = Diversity(self.div_model, self.div_tokenizer, device_div)
+        biscope = BiScope(self.bi_model, self.bi_tokenizer, self.device_bi)
+
+        diveye_features = diveye.compute_features(text)
+        biscope_features = biscope.detect_single_sample(text)
 
         for f in biscope_features:
             diveye_features.append(f)