Update app.py

app.py

@@ -2,76 +2,55 @@ import gc
 import logging
 import os
 import re
-
-
-import spaces
-
 import torch
 from cleantext import clean
 import gradio as gr
 from tqdm.auto import tqdm
 from transformers import pipeline
-from transformers import AutoModelForSequenceClassification, AutoTokenizer
 
 logging.basicConfig(level=logging.INFO)
 logging.info(f"torch version:\t{torch.__version__}")
 
+# --- 1. ต้องประกาศชื่อ Model ไว้ตรงนี้ก่อน (ห้ามย้ายไปไว้ข้างล่าง) ---
+checker_model_name = "textattack/roberta-base-CoLA"
+corrector_model_name = "pszemraj/flan-t5-large-grammar-synthesis"
+
+# --- 2. เช็ค Device (ป้องกัน RuntimeError เรื่อง NVIDIA) ---
 device = 0 if torch.cuda.is_available() else -1
+logging.info(f"Using device: {'cuda' if device == 0 else 'cpu'}")
 
+# --- 3. สร้าง Pipeline (ดึงตัวแปรจากข้อ 1 มาใช้) ---
 checker = pipeline(
     "text-classification",
     model=checker_model_name,
-    device=device, # แก้จาก device_map="cuda" เป็น device
+    device=device,
 )
-
 corrector = pipeline(
     "text2text-generation",
     model=corrector_model_name,
-    device=device, # แก้จาก device_map="cuda" เป็น device
-)
-
-corrector = pipeline(
-    "text2text-generation",
-    corrector_model_name,
-    device_map="cuda",
+    device=device,
 )
 
+# --- ฟังก์ชันการทำงานอื่นๆ ---
 def split_text(text: str) -> list:
-    # Split the text into sentences using regex
     sentences = re.split(r"(?<=[^A-Z].[.?]) +(?=[A-Z])", text)
-
-    # Initialize lists for batching
     sentence_batches = []
     temp_batch = []
-
-    # Create batches of 2-3 sentences
     for sentence in sentences:
         temp_batch.append(sentence)
-        if len(temp_batch) >= 2 and len(temp_batch) <= 3 or sentence == sentences[-1]:
+        if (len(temp_batch) >= 2 and len(temp_batch) <= 3) or sentence == sentences[-1]:
             sentence_batches.append(temp_batch)
             temp_batch = []
-
     return sentence_batches
 
-
 def correct_text(text: str, separator: str = " ") -> str:
-
-    # Split the text into sentence batches
     sentence_batches = split_text(text)
-
-    # Initialize a list to store the corrected text
     corrected_text = []
-
-    # Process each batch
-    for batch in tqdm(
-        sentence_batches, total=len(sentence_batches), desc="correcting text.."
-    ):
+    for batch in tqdm(sentence_batches, desc="correcting text.."):
         raw_text = " ".join(batch)
-
-        # Check grammar quality
         results = checker(raw_text)
-
-        # Correct text if needed
+        
+        # ตรวจสอบคุณภาพไวยากรณ์
         if results[0]["label"] != "LABEL_1" or (
             results[0]["label"] == "LABEL_1" and results[0]["score"] < 0.9
         ):
@@ -79,42 +58,19 @@ def correct_text(text: str, separator: str = " ") -> str:
             corrected_text.append(corrected_batch[0]["generated_text"])
         else:
             corrected_text.append(raw_text)
-
-    # Join the corrected text
     return separator.join(corrected_text)
 
-
 def update(text: str):
-    # Clean and truncate input text
     text = clean(text[:4000], lower=False)
     return correct_text(text)
 
-
-# Create the Gradio interface
+# --- 4. Interface ---
 with gr.Blocks() as demo:
-    gr.Markdown("# <center>Robust Grammar Correction with FLAN-T5</center>")
-    gr.Markdown(
-        "**Instructions:** Enter the text you want to correct in the textbox below (_text will be truncated to 4000 characters_). Click 'Process' to run."
-    )
-    gr.Markdown(
-        """Models:
-    - `textattack/roberta-base-CoLA` for grammar quality detection
-    - `pszemraj/flan-t5-large-grammar-synthesis` for grammar correction
-    """
-    )
+    gr.Markdown("# <center>Robust Grammar Correction</center>")
     with gr.Row():
-        inp = gr.Textbox(
-            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)
+        inp = gr.Textbox(label="Input", placeholder="Enter text here...")
+        out = gr.Textbox(label="Output", interactive=False)
     btn = gr.Button("Process")
     btn.click(fn=update, inputs=inp, outputs=out)
-    gr.Markdown("---")
-    gr.Markdown(
-        "- See the [model card](https://huggingface.co/pszemraj/flan-t5-large-grammar-synthesis) for more info"
-    )
 
-# Launch the demo
-demo.launch(debug=True)
+demo.launch()