Update app.py

app.py

@@ -10,17 +10,20 @@ import string
 nltk.download('punkt')
 nltk.download('stopwords')
 
+# Check for GPU and set the device accordingly
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
 # Load AI Detector model and tokenizer from Hugging Face (DistilBERT)
 tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
-model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
+model = AutoModelForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english").to(device)
 
 # Load SRDdev Paraphrase model and tokenizer for humanizing text
 paraphrase_tokenizer = T5Tokenizer.from_pretrained("SRDdev/Paraphrase")
-paraphrase_model = T5ForConditionalGeneration.from_pretrained("SRDdev/Paraphrase")
+paraphrase_model = T5ForConditionalGeneration.from_pretrained("SRDdev/Paraphrase").to(device)
 
 # AI detection function using DistilBERT
 def detect_ai_generated(text):
-    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
+    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512).to(device)
     with torch.no_grad():
         outputs = model(**inputs)
     probabilities = torch.softmax(outputs.logits, dim=1)
@@ -98,7 +101,7 @@ def humanize_text(AI_text):
     paraphrased_paragraphs = []
     for paragraph in paragraphs:
         if paragraph.strip():
-            inputs = paraphrase_tokenizer(paragraph, return_tensors="pt", max_length=512, truncation=True)
+            inputs = paraphrase_tokenizer(paragraph, return_tensors="pt", max_length=512, truncation=True).to(device)
             paraphrased_ids = paraphrase_model.generate(
                 inputs['input_ids'],
                 max_length=inputs['input_ids'].shape[-1] + 20,  # Slightly more than the original input length