Update app.py

app.py

@@ -1,80 +1,79 @@
-import os
-import time
 import streamlit as st
-import google.generativeai as genai
-from google.auth import default
-from dotenv import load_dotenv
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+from spellchecker import SpellChecker
+import re
+import torch
 
-# ✅ Set page config FIRST
-st.set_page_config(page_title="Grammar & Spelling Assistant", page_icon="🧠")
+# Load model and tokenizer
+@st.cache_resource
+def load_model():
+    model_name = "vennify/t5-base-grammar-correction"
+    tokenizer = T5Tokenizer.from_pretrained(model_name)
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
+    return tokenizer, model
 
-# Load environment variables
-load_dotenv()
+tokenizer, model = load_model()
 
-# ✅ Use Application Default Credentials
-if os.getenv("GOOGLE_APPLICATION_CREDENTIALS"):
-    credentials, _ = default()
-    genai.configure(credentials=credentials)
-else:
-    st.error("❌ GOOGLE_APPLICATION_CREDENTIALS is not set. Please add it to .env or Hugging Face Secrets.")
-    st.stop()
+# Step 1: Fix typos using pyspellchecker
+def correct_spelling(text):
+    spell = SpellChecker()
+    words = re.findall(r'\b\w+\b|\S', text)
+    corrected_words = []
 
-# Initialize Gemini model
-model = genai.GenerativeModel("gemini-2.0-pro")
+    for word in words:
+        # 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)
 
-# Optional styling
-st.markdown("""
-<style>
-body, .stApp { background-color: #121212 !important; color: #e0e0e0 !important; }
-.stChatInput { background: #222 !important; border: 1px solid #555 !important; }
-</style>
-""", unsafe_allow_html=True)
+    return ' '.join(corrected_words)
 
-# App header
-st.title("📚 Grammar Guardian")
-st.caption("Correct grammar, get explanations, and improve your writing!")
+# 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)
+    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
 
-# Initialize chat history
-if "history" not in st.session_state:
-    st.session_state.history = []
+# 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.")
 
-# Display chat history
-for message in st.session_state.history:
-    role, content = message["role"], message["content"]
-    with st.chat_message(role):
-        st.markdown(content)
+user_input = st.text_area("✍️ Enter your sentence below:", height=150)
 
-# Chat input
-prompt = st.chat_input("Type a sentence you'd like to improve...")
+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)
 
-if prompt:
-    # Show user message
-    with st.chat_message("user"):
-        st.markdown(prompt)
-    st.session_state.history.append({"role": "user", "content": prompt})
+        # Step 2: Grammar correction
+        final_output = correct_grammar(spelling_fixed)
 
-    # Process input with Gemini
-    with st.spinner("Analyzing..."):
-        try:
-            full_prompt = f"""
-You are a grammar correction assistant.
-When a user gives a sentence, do two things:
-1. Correct the sentence.
-2. Explain clearly why you corrected it.
+        st.markdown("### ✅ Final Correction:")
+        st.success(final_output)
 
-Respond using this format:
-**Correction:** <Corrected sentence>
-**Explanation:** <Explanation of why it was corrected>
+        st.markdown("### 🔍 Explanation of Changes:")
+        st.info(f"""
+**Original Sentence:**  
+{user_input}
 
-Sentence: {prompt}
-"""
-            response = model.generate_content(full_prompt)
-            result = response.text
+**After Spelling Correction:**  
+{spelling_fixed}
 
-            # Show assistant message
-            with st.chat_message("assistant"):
-                st.markdown(result)
-            st.session_state.history.append({"role": "assistant", "content": result})
+**After Grammar Correction:**  
+{final_output}
 
-        except Exception as e:
-            st.error(f"Error from Gemini API: {e}")
+**Explanation:**  
+- Typos like `ober` → `over`, `dogz#` → `dogs` were corrected.  
+- Grammar and punctuation (like capitalization) were fixed.  
+- Your original message and word choices were preserved.
+""")