DivEye - PR (fixes x2)

app.py

@@ -1,26 +1,20 @@
 import gradio as gr
 import numpy as np
+import os
 import pandas as pd
 from software import Software
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 
-software = None
 theme = gr.Theme.from_hub("gstaff/xkcd") 
 
-def warmup():
-    global software
-    print("[DivEye] Warming up models...")
-    software = Software()
-    print("[DivEye] Models are ready.")
-
 def detect_ai_text(text):
-    global software
     if software is None:
-        warmup()
-        return "⏳ Models not loaded yet. Please wait for a few moments.", 0.0, pd.DataFrame({
-                "Source": ["AI Generated", "Human Written"],
-                "Probability (%)": [0, 0]
-            })
-
+        return "❗ Model not loaded. We require a GPU to run DivEye.", 0.0, pd.DataFrame({
+            "Source": ["AI Generated", "Human Written"],
+            "Probability (%)": [0, 0]
+        })
+    
     text = text.strip()
     if not text or len(text.split()) < 15:
         return (
@@ -50,6 +44,29 @@ def detect_ai_text(text):
 
     return message, round(ai_prob, 3), bar_data
 
+# Token from environment variable
+token = os.getenv("HF_TOKEN") 
+
+if not torch.cuda.is_available():
+    print("[DivEye] CUDA not available. Running on CPU.")
+    DESCRIPTION = "This demo requires a GPU to run efficiently. Please use a machine with CUDA support."
+
+# Import necessary models and tokenizers
+if torch.cuda.is_available():
+    model_name_div = "tiiuae/falcon-7b"
+    model_name_bi = "google/gemma-1.1-2b-it"
+
+    div_model = AutoModelForCausalLM.from_pretrained(model_name_div, torch_dtype=torch.float16, device_map="cuda:0", use_auth_token=token)
+    div_tokenizer = AutoTokenizer.from_pretrained(model_name_div, use_fast=False, trust_remote_code=True, use_auth_token=token)
+
+    bi_model = AutoModelForCausalLM.from_pretrained(model_name_bi, torch_dtype=torch.float16, device_map="cuda:1", use_auth_token=token)
+    bi_tokenizer = AutoTokenizer.from_pretrained(model_name_bi, use_fast=False, trust_remote_code=True, use_auth_token=token)
+
+    div_model.eval()
+    bi_model.eval()
+
+    software = Software(div_model, div_tokenizer, bi_model, bi_tokenizer, div_model.device, bi_model.device)
+
 # Gradio app setup
 with gr.Blocks(title="DivEye") as demo:
     gr.HTML("""

software.py

@@ -15,7 +15,7 @@ import os
 class Diversity:
     def __init__(self, model, tokenizer, device):
         self.tokenizer = tokenizer
-        self.model = model.to(device)
+        self.model = model
         self.device = device
         
     def compute_log_likelihoods(self, text):
@@ -56,7 +56,7 @@ class BiScope:
     def __init__(self, model, tokenizer, device):
         self.COMPLETION_PROMPT_ONLY = "Complete the following text: "
         self.tokenizer = tokenizer
-        self.model = model.to(device)
+        self.model = model
         self.device = device
     
     def compute_fce_loss(self, logits, targets, text_slice):
@@ -97,70 +97,23 @@ class BiScope:
 
 
 class Software:
-    def __init__(self):
-        self.token = os.getenv("HF_TOKEN") 
-        self.device_div = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-        self.device_bi = self.device_div
+    def __init__(self, div_model, div_tokenizer, bi_model, bi_tokenizer, device_div="cuda", device_bi="cuda"):
+        self.div_model = div_model
+        self.div_tokenizer = div_tokenizer
+        self.bi_model = bi_model
+        self.bi_tokenizer = bi_tokenizer 
 
-        self.div_model = None
-        self.div_tokenizer = None
-        self.bi_model = None
-        self.bi_tokenizer = None        
+        self.device_div = device_div
+        self.device_bi = device_bi    
         
         self.model_path = Path(__file__).parent / "model.json"
 
         self.model = xgb.XGBClassifier()
         self.model.load_model(self.model_path)
-
-    def _load_div_models(self):
-        if self.div_model is None or self.div_tokenizer is None:
-            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="cuda",
-                torch_dtype=torch.float16,
-                trust_remote_code=True,
-                use_auth_token=self.token
-            )
-            self.div_model.to(self.device_div)
-
-    def _load_bi_models(self):
-        if self.bi_model is None or self.bi_tokenizer is None:
-            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="cuda",
-                torch_dtype=torch.float16,
-                trust_remote_code=True,
-                use_auth_token=self.token
-            )
-            self.bi_model.to(self.device_bi)
-        
-    def load_data(self, jsonl_path):
-        ids, texts = [], []
-        with open(jsonl_path, 'r') as f:
-            for line in f:
-                obj = json.loads(line)
-                ids.append(obj["id"])
-                texts.append(obj["text"])
-        return ids, texts
     
     @spaces.GPU
-    def evaluate(self, text):
-        self._load_div_models()
-        self._load_bi_models()
-        
-        # 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)
+    def evaluate(self, text):     
+        diveye = Diversity(self.div_model, self.div_tokenizer, self.device_div)
         biscope = BiScope(self.bi_model, self.bi_tokenizer, self.device_bi)
 
         diveye_features = diveye.compute_features(text)