Update app.py

app.py

@@ -7,6 +7,8 @@ from huggingface_hub import login, HfFolder
 from transformers import AutoTokenizer, AutoModelForSequenceClassification
 from scipy.special import softmax
 import logging
+import spaces
+
 
 # Setup logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(message)s')
@@ -62,7 +64,8 @@ models = {path: AutoModelForSequenceClassification.from_pretrained(path) for pat
 
 def get_quality_name(model_name):
     return quality_mapping.get(model_name.split('/')[-1], "Unknown Quality")
-
+    
+@spaces.GPU
 def model_prediction(model, text, device):
     model.to(device)
     model.eval()
@@ -91,7 +94,7 @@ def main_interface(text):
     for model_path, model in models.items():
         quality_name = get_quality_name(model_path)
         avg_prob = model_prediction(model, text, device)
-        if avg_prob >= 0.995:  # Only consider probabilities >= 0.90
+        if avg_prob >= 0.95:  # Only consider probabilities >= 0.90
             results.append((quality_name, avg_prob))
         logging.info(f"Model: {model_path}, Quality: {quality_name}, Average Probability: {avg_prob:.3f}")