updates to app.py

app.py

@@ -3,12 +3,20 @@
 from typing import List
 from transformers import pipeline
 import os
+import torch
 
 app = FastAPI()
-# Initialize the NSFW model once
+
 env = os.environ.get("env", "dev")
-torch_device = "cuda" if env == "prod" else "cpu" 
 
+if env == "prod" and torch.cuda.is_available():
+    torch_device = "cuda"
+elif env == "prod" and not torch.cuda.is_available():
+    print("Warning: CUDA is not available. Falling back to CPU.")
+    torch_device = "cpu"
+else:
+    torch_device = "cpu"
+
 class InferenceData(BaseModel):
     name: str
     shape: List[int]

validator/main.py

@@ -98,7 +98,7 @@ def validate_each_sentence(
         # Split the value into sentences using nltk sentence tokenizer.
         sentences = nltk.sent_tokenize(value)
 
-        unsupported_sentences, supported_sentences = [], []
+        clean_sentences, nsfw_sentences = [], []
         error_spans: List[ErrorSpan] = []
         char_index = 0
 
@@ -108,7 +108,7 @@ def validate_each_sentence(
             pred_labels = sentence_predictions[idx]
 
             if pred_labels:
-                unsupported_sentences.append(sentence)
+                nsfw_sentences.append(sentence)
                 error_spans.append(
                     ErrorSpan(
                         start=char_index,
@@ -117,11 +117,11 @@ def validate_each_sentence(
                     )
                 )
             else:
-                supported_sentences.append(sentence)
+                clean_sentences.append(sentence)
             char_index += len(sentence) + 1  
 
-        if unsupported_sentences:
-            unsupported_sentences_text = "- " + "\n- ".join(unsupported_sentences)
+        if nsfw_sentences:
+            nsfw_sentences_text = "- " + "\n- ".join(nsfw_sentences)
 
             return FailResult(
                 metadata=metadata,