Update app.py

app.py

@@ -3,6 +3,7 @@ import torch
 import torch.nn as nn
 from transformers import AutoTokenizer
 from huggingface_hub import hf_hub_download
+import spaces
 
 # 1. Model Architecture
 class SourceCodeAuthorCheck(nn.Module):
@@ -32,22 +33,24 @@ class SourceCodeAuthorCheck(nn.Module):
         return self.fc(pooled)
 
 # 2. Device and Loading Initialization
-# Note: HF Free Spaces use CPU. We check for CUDA in case you upgrade the Space hardware.
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 tokenizer = AutoTokenizer.from_pretrained("gpt2")
 tokenizer.pad_token = tokenizer.eos_token
 
-model = SourceCodeAuthorCheck().to(device)
-
-# Download weights securely from your model repository
+# Load model globally on CPU first
+model = SourceCodeAuthorCheck()
 model_path = hf_hub_download(repo_id="assix-research/SourceCodeAuthorCheck-SLM-10M", filename="source_code_classifier.pth")
-model.load_state_dict(torch.load(model_path, map_location=device, weights_only=True))
+model.load_state_dict(torch.load(model_path, map_location="cpu", weights_only=True))
 model.eval()
 
-# 3. Inference Logic
+# 3. Inference Logic with ZeroGPU Decorator
+@spaces.GPU
 def predict_author(code_snippet):
     if not code_snippet or not code_snippet.strip():
         return "Please paste valid code.", "0.0%"
+    
+    # ZeroGPU dynamically provides CUDA access inside this decorated function
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    model.to(device)
         
     inputs = tokenizer(
         code_snippet, 
@@ -58,7 +61,6 @@ def predict_author(code_snippet):
     ).to(device)
     
     with torch.no_grad():
-        # Handle mixed precision safely based on available hardware
         if torch.cuda.is_available():
             with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
                 logits = model(inputs['input_ids'], inputs['attention_mask'])
@@ -70,6 +72,9 @@ def predict_author(code_snippet):
     score = round(prob * 100, 2)
     verdict = "🤖 AI Generated" if prob > 0.5 else "👨‍💻 Human Written"
     
+    # Move model back to CPU to free up ZeroGPU vRAM for other users
+    model.to("cpu")
+    
     return verdict, f"{score}%"
 
 # 4. Gradio Interface Construction