Update app.py

app.py

@@ -16,15 +16,18 @@ from transformers import AutoModelForSequenceClassification, AutoTokenizer
 logging.basicConfig(level=logging.INFO)
 logging.info(f"torch version:\t{torch.__version__}")
 
-# Model names
-checker_model_name = "textattack/roberta-base-CoLA"
-corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
-
+device = 0 if torch.cuda.is_available() else -1
 
 checker = pipeline(
     "text-classification",
-    checker_model_name,
-    device_map="cuda",
+    model=checker_model_name,
+    device=device, # แก้จาก device_map="cuda" เป็น device
+)
+
+corrector = pipeline(
+    "text2text-generation",
+    model=corrector_model_name,
+    device=device, # แก้จาก device_map="cuda" เป็น device
 )
 
 corrector = pipeline(