Update app.py

app.py

@@ -1,22 +1,49 @@
 import streamlit as st
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, pipeline
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, T5Tokenizer, T5ForConditionalGeneration
 
-# Load grammar correction model
-model_name = "vennify/t5-base-grammar-correction"
-tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
-pipe = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+# Model 1: Grammar & spelling correction
+grammar_model = "vennify/t5-base-grammar-correction"
+tokenizer = AutoTokenizer.from_pretrained(grammar_model)
+model = AutoModelForSeq2SeqLM.from_pretrained(grammar_model)
+corrector = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-st.title("✍️ English Grammar & Writing Assistant")
+# Model 2: Explanation & suggestions
+explain_model = "google/flan-t5-large"
+exp_tokenizer = T5Tokenizer.from_pretrained(explain_model)
+exp_model = T5ForConditionalGeneration.from_pretrained(explain_model)
 
-text = st.text_area("Enter your sentence or paragraph:")
+def correct_text(text):
+    input_text = "gec: " + text
+    result = corrector(input_text, max_length=512, clean_up_tokenization_spaces=True)
+    return result[0]['generated_text']
 
-if st.button("Correct Text"):
-    if not text.strip():
-        st.warning("Please enter some text.")
+def explain_and_suggest(original, corrected):
+    prompt = (
+        f"Original: {original}\n"
+        f"Corrected: {corrected}\n"
+        "Explain the grammar, spelling, or style changes made. "
+        "Mention any alternate better words. Then suggest how the writer can improve in the future."
+    )
+    inputs = exp_tokenizer(prompt, return_tensors="pt", truncation=True)
+    outputs = exp_model.generate(**inputs, max_length=512)
+    return exp_tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+# Streamlit UI
+st.set_page_config(page_title="Grammar & Writing Coach", layout="centered")
+st.title("📝 Grammar & Writing Coach")
+st.markdown("Enter a sentence, paragraph, or essay. We'll correct it, explain the changes, and help you improve.")
+
+text_input = st.text_area("✍️ Your text here:", height=200)
+
+if st.button("🔍 Analyze & Improve"):
+    if not text_input.strip():
+        st.warning("Please enter some text first.")
     else:
-        input_text = "gec: " + text
-        result = pipe(input_text, max_length=512)
-        corrected = result[0]['generated_text']
-        st.subheader("✅ Corrected Text:")
+        corrected = correct_text(text_input)
+        explanation = explain_and_suggest(text_input, corrected)
+
+        st.subheader("✅ Corrected Version:")
         st.write(corrected)
+
+        st.subheader("📘 Explanation & Suggestions:")
+        st.write(explanation)