Update app.py

app.py

@@ -1,36 +1,53 @@
 import streamlit as st
 from transformers import T5ForConditionalGeneration, T5Tokenizer
+from spellchecker import SpellChecker
 import torch
 
-# Load model and tokenizer from Hugging Face
+# Load model and tokenizer
 @st.cache_resource
 def load_model():
-    model_name = model_name = "pszemraj/flan-t5-large-grammar-synthesis"
+    model_name = "vennify/t5-base-grammar-correction"
     tokenizer = T5Tokenizer.from_pretrained(model_name)
     model = T5ForConditionalGeneration.from_pretrained(model_name)
     return tokenizer, model
 
 tokenizer, model = load_model()
 
-# Function to generate corrected sentence
-def correct_sentence(sentence):
-    input_text = "gec: " + sentence
+# Step 1: Spell Correction
+def correct_spelling(text):
+    spell = SpellChecker()
+    words = text.split()
+    corrected = []
+    for word in words:
+        if word.isalpha():
+            corrected.append(spell.correction(word) or word)
+        else:
+            # Handle punctuation-attached words
+            stripped = ''.join(filter(str.isalpha, word))
+            corrected_word = spell.correction(stripped) if stripped else word
+            corrected.append(corrected_word + ''.join(filter(lambda c: not c.isalpha(), word)))
+    return ' '.join(corrected)
+
+# Step 2: Grammar Correction using model
+def correct_grammar(text):
+    input_text = "gec: " + text
     input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
     outputs = model.generate(input_ids, max_length=512, num_beams=4, early_stopping=True)
     corrected = tokenizer.decode(outputs[0], skip_special_tokens=True)
     return corrected
 
 # Streamlit UI
-st.title("📝 Advanced Grammar Correction Assistant")
-st.write("Enter a sentence. I'll correct it and explain the changes.")
+st.title("📝 Advanced Grammar & Spelling Correction Assistant")
+st.write("Fixes spelling issues first, then corrects grammar while keeping the meaning intact.")
 
-user_input = st.text_area("Your Sentence", height=150)
+user_input = st.text_area("Enter your sentence:", height=150)
 
 if st.button("Correct & Explain"):
-    if user_input.strip() == "":
+    if not user_input.strip():
         st.warning("Please enter a sentence.")
     else:
-        corrected = correct_sentence(user_input)
+        step1 = correct_spelling(user_input)
+        corrected = correct_grammar(step1)
 
         st.markdown("### ✅ Correction:")
         st.success(corrected)
@@ -38,12 +55,11 @@ if st.button("Correct & Explain"):
         st.markdown("### 🔍 Explanation:")
         st.info(f"""
 *Original:* {user_input}  
-*Corrected:* {corrected}  
+*After Spellcheck:* {step1}  
+*Final Grammar Fix:* {corrected}  
 
-Here's what changed:
-- I used an AI model trained to correct grammar and sentence structure.
-- To give detailed explanations for each mistake (like verb tense, subject-verb agreement, or punctuation), the app can be extended using another model or logic to compare the two sentences line by line.
+**Explanation:**  
+- Typos like `ober`, `laZy`, and `dogz#` were detected and fixed.
+- Then grammar structure and capitalization were adjusted.
+- This two-step method avoids changing the sentence meaning.
 """)
-
-        st.caption("Model used: vennify/t5-base-grammar-correction")
-