DivEye - PR (fixes)

software.py

@@ -9,6 +9,7 @@ from scipy.stats import skew, kurtosis, entropy
 from tqdm import tqdm
 from torch.nn import CrossEntropyLoss
 from pathlib import Path
+import spaces
 
 class Diversity:
     def __init__(self, model, tokenizer, device):
@@ -102,14 +103,14 @@ class Software:
         else:
             self.device_bi = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
 
-        self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True, local_files_only=True)
+        self.div_tokenizer = AutoTokenizer.from_pretrained("tiiuae/falcon-7b", use_fast=False, trust_remote_code=True)
         self.div_model = AutoModelForCausalLM.from_pretrained(
-            "tiiuae/falcon-7b", device_map=self.device_div, torch_dtype=torch.float16, trust_remote_code=True, local_files_only=True
+            "tiiuae/falcon-7b", device_map=self.device_div, torch_dtype=torch.float16, trust_remote_code=True
         )
         
-        self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True, local_files_only=True)
+        self.bi_tokenizer = AutoTokenizer.from_pretrained("google/gemma-1.1-2b-it", use_fast=False, trust_remote_code=True)
         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, local_files_only=True
+            "google/gemma-1.1-2b-it", device_map=self.device_bi, torch_dtype=torch.float16, trust_remote_code=True
         )
         
         self.diveye = Diversity(self.div_model, self.div_tokenizer, self.device_div)
@@ -128,6 +129,7 @@ class Software:
                 texts.append(obj["text"])
         return ids, texts
     
+    @spaces.GPU
     def evaluate(self, text):
         diveye_features = self.diveye.compute_features(text)
         biscope_features = self.biscope.detect_single_sample(text)