Spaces:
Running
Running
import streamlit as st | |
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer | |
import torch | |
st.set_page_config(page_title="Grammar Corrector", page_icon="π") | |
st.title("π AI Grammar & Spell Corrector") | |
def load_model(): | |
model = AutoModelForSeq2SeqLM.from_pretrained("vennify/t5-base-grammar-correction") | |
tokenizer = AutoTokenizer.from_pretrained("vennify/t5-base-grammar-correction") | |
return model, tokenizer | |
model, tokenizer = load_model() | |
def correct_grammar(text): | |
input_text = "grammar: " + text | |
inputs = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True) | |
with torch.no_grad(): | |
outputs = model.generate(inputs, max_length=512, num_beams=4, early_stopping=True) | |
corrected = tokenizer.decode(outputs[0], skip_special_tokens=True) | |
return corrected | |
text = st.text_area("Enter your text:", height=200) | |
if st.button("Correct Text"): | |
if text.strip(): | |
with st.spinner("Correcting..."): | |
corrected = correct_grammar(text) | |
st.subheader("β Corrected Text") | |
st.write(corrected) | |
else: | |
st.warning("Please enter some text.") | |