Update app.py

app.py

@@ -3,9 +3,7 @@ import logging
 import os
 import re
 
-
 import spaces
-
 import torch
 from cleantext import clean
 import gradio as gr
@@ -16,21 +14,25 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 logging.basicConfig(level=logging.INFO)
 logging.info(f"torch version:\t{torch.__version__}")
 
+# Ensure compatibility with GPU or CPU
+device = "cuda" if torch.cuda.is_available() else "cpu"
+logging.info(f"Using device: {device}")
+
 # Model names
 checker_model_name = "textattack/roberta-base-CoLA"
 corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
 
-
+# Initialize pipelines with device mapping
 checker = pipeline(
     "text-classification",
     checker_model_name,
-    device_map="cuda",
+    device=0 if torch.cuda.is_available() else -1,  # 0 for GPU, -1 for CPU
 )
 
 corrector = pipeline(
     "text2text-generation",
     corrector_model_name,
-    device_map="cuda",
+    device=0 if torch.cuda.is_available() else -1,  # 0 for GPU, -1 for CPU
 )
 
 def split_text(text: str) -> list:
@@ -53,7 +55,6 @@ def split_text(text: str) -> list:
 
 @spaces.GPU(duration=60)
 def correct_text(text: str, separator: str = " ") -> str:
-
     # Split the text into sentence batches
     sentence_batches = split_text(text)
 
@@ -102,11 +103,11 @@ with gr.Blocks() as demo:
     )
     with gr.Row():
         inp = gr.Textbox(
-            label="input",
+            label="Input",
             placeholder="Enter text to check & correct",
             value="I wen to the store yesturday to bye some food. I needd milk, bread, and a few otter things. The store was really crowed and I had a hard time finding everyting I needed. I finaly made it to the check out line and payed for my stuff.",
         )
-        out = gr.Textbox(label="output", interactive=False)
+        out = gr.Textbox(label="Output", interactive=False)
     btn = gr.Button("Process")
     btn.click(fn=update, inputs=inp, outputs=out)
     gr.Markdown("---")
@@ -115,4 +116,4 @@ with gr.Blocks() as demo:
     )
 
 # Launch the demo
-demo.launch(debug=True)
+demo.launch(debug=True)