Spaces:
Sleeping
Sleeping
import streamlit as st | |
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
import difflib | |
# Load model and tokenizer | |
def load_model(): | |
model_name = "vennify/t5-base-grammar-correction" | |
tokenizer = AutoTokenizer.from_pretrained(model_name) | |
model = AutoModelForSeq2SeqLM.from_pretrained(model_name) | |
return tokenizer, model | |
tokenizer, model = load_model() | |
# Function to correct grammar | |
def correct_text(text): | |
input_text = "gec: " + text | |
inputs = tokenizer.encode(input_text, return_tensors="pt", truncation=True) | |
outputs = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True) | |
corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return corrected_text | |
# Word-by-word diff | |
def get_word_diffs(original, corrected): | |
orig_words = original.split() | |
corr_words = corrected.split() | |
diff = list(difflib.ndiff(orig_words, corr_words)) | |
return diff | |
# Streamlit UI | |
st.set_page_config(page_title="Grammar Agent", layout="wide") | |
st.title("π Grammar Agent β Like Grammarly, Powered by AI") | |
st.markdown(""" | |
This AI grammar agent helps you: | |
- β Correct grammar, spelling, and punctuation | |
- β Improve sentence clarity | |
- β See word-by-word differences | |
- β Get suggestions on how to rewrite | |
""") | |
text_input = st.text_area("Enter your sentence, paragraph, or essay:", height=200) | |
if st.button("Correct Grammar") and text_input.strip(): | |
corrected = correct_text(text_input) | |
diffs = get_word_diffs(text_input, corrected) | |
st.subheader("π§ Corrected Version") | |
st.success(corrected) | |
st.subheader("π Word-by-Word Differences") | |
diff_text = [] | |
for word in diffs: | |
if word.startswith("-"): | |
diff_text.append(f"β {word[2:]}") | |
elif word.startswith("+"): | |
diff_text.append(f"β {word[2:]}") | |
st.markdown(" ".join(diff_text)) | |
st.subheader("π‘ Writing Suggestion") | |
if len(text_input.split()) > len(corrected.split()): | |
st.info("Try simplifying your sentences for better clarity.") | |
elif len(corrected.split()) > len(text_input.split()): | |
st.info("The AI added helpful detail β good for formal writing.") | |
else: | |
st.info("The AI cleaned grammar and punctuation with minimal changes.") | |