Update app.py

app.py

@@ -14,20 +14,26 @@ def load_model():
 
 tokenizer, model = load_model()
 
-# Step 1: Fix obvious typos using spellchecker
+# Step 1: Fix typos using pyspellchecker
 def correct_spelling(text):
     spell = SpellChecker()
     words = re.findall(r'\b\w+\b|\S', text)
     corrected_words = []
 
     for word in words:
-        if word.isalpha():
-            corrected_words.append(spell.correction(word.lower()) or word)
+        # Remove non-alphanumeric characters for spellcheck
+        clean_word = re.sub(r'[^\w\s]', '', word)
+        if clean_word.isalpha():
+            corrected_word = spell.correction(clean_word.lower()) or clean_word
+            # Restore punctuation
+            trailing = ''.join(re.findall(r'[^\w\s]', word))
+            corrected_words.append(corrected_word + trailing)
         else:
             corrected_words.append(word)
+
     return ' '.join(corrected_words)
 
-# Step 2: Run grammar correction model
+# Step 2: Grammar Correction using the model
 def correct_grammar(text):
     input_text = "gec: " + text
     input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
@@ -35,35 +41,39 @@ def correct_grammar(text):
     corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected
 
-# Streamlit UI
-st.title("🧠 AI Grammar & Spelling Fixer (No Weird Word Changes)")
-st.write("Fixes grammar and spelling without changing your intended words.")
+# UI
+st.set_page_config(page_title="Grammar & Spelling Assistant", page_icon="🧠")
+st.title("🧠 Grammar & Spelling Correction Assistant")
+st.write("Fixes your typos and grammar without changing your meaning.")
 
-user_input = st.text_area("✍️ Type your sentence:", height=150)
+user_input = st.text_area("✍️ Enter your sentence below:", height=150)
 
 if st.button("Correct & Explain"):
     if not user_input.strip():
         st.warning("Please enter a sentence.")
     else:
+        # Step 1: Spelling correction
         spelling_fixed = correct_spelling(user_input)
+
+        # Step 2: Grammar correction
         final_output = correct_grammar(spelling_fixed)
 
-        st.markdown("### ✅ Correction:")
+        st.markdown("### ✅ Final Correction:")
         st.success(final_output)
 
-        st.markdown("### 🔍 Explanation:")
+        st.markdown("### 🔍 Explanation of Changes:")
         st.info(f"""
-**Original Input:**  
+**Original Sentence:**  
 {user_input}
 
-**After Spell Check (typo fix only):**  
+**After Spelling Correction:**  
 {spelling_fixed}
 
 **After Grammar Correction:**  
 {final_output}
 
-**What Changed:**  
-- Fixed typos: `ober` → `over`, `dogz#` → `dogs#`
-- Grammar and capitalization improved
-- Meaning and keywords like "brown fox" and "lazy dog" were preserved
+**Explanation:**  
+- Typos like `ober` → `over`, `dogz#` → `dogs` were corrected.  
+- Grammar and punctuation (like capitalization) were fixed.  
+- Your original message and word choices were preserved.
 """)