Update app.py

app.py

@@ -1,46 +1,46 @@
 import streamlit as st
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-import torch
 
-# Set up the Streamlit app
-st.title("Correct your Grammar with Transformers")
+st.title("Correct Grammar with Transformers ")
 st.write("")
 st.write("Input your text here!")
 
-# Create input text area
 default_value = "Mike and Anna is skiing"
 sent = st.text_area("Text", default_value, height=50)
+num_return_sequences = st.sidebar.number_input('Number of Return Sequences', min_value=1, max_value=3, value=1, step=1)
+
+# Run Model
+from transformers import T5ForConditionalGeneration, T5Tokenizer
+import torch
+
+torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
+tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')
+model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector').to(torch_device)
 
-# Create "Check Now" button
+def correct_grammar(input_text, num_return_sequences=num_return_sequences):
+ batch = tokenizer([input_text], truncation=True, padding='max_length', max_length=len(input_text), return_tensors="pt").to(torch_device)
+ results = model.generate(**batch, max_length=len(input_text), num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)
+ return results
+
+# Prompts
+results = correct_grammar(sent, num_return_sequences)
+
+# Decode generated sequences
+generated_sequences = [tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True, skip_special_tokens=True) for generated_sequence in results]
+
+# Add "Check Now" button
 if st.button("Check Now"):
- # Run Model
- torch_device = 'cuda' if torch.cuda.is_available() else 'cpu'
- tokenizer = T5Tokenizer.from_pretrained('deep-learning-analytics/GrammarCorrector')
- model = T5ForConditionalGeneration.from_pretrained('deep-learning-analytics/GrammarCorrector').to(torch_device)
-
- def correct_grammar(input_text, num_return_sequences=1):
- batch = tokenizer([input_text], truncation=True, padding='max_length', max_length=len(input_text), return_tensors="pt").to(torch_device)
- results = model.generate(**batch, max_length=len(input_text), num_beams=2, num_return_sequences=num_return_sequences, temperature=1.5)
- return results
-
- # Prompts
- results = correct_grammar(sent, num_return_sequences=1)
-
- # Decode results
- generated_sequences = []
- for generated_sequence_idx, generated_sequence in enumerate(results):
- text = tokenizer.decode(generated_sequence, clean_up_tokenization_spaces=True, skip_special_tokens=True)
- generated_sequences.append(text)
-
- # Check correctness
- is_correct = sent == generated_sequences[0]
-
- # Display correctness result
- if is_correct:
- st.write("Result: ", generated_sequences[0], " (Correct)", key="result_text", unsafe_allow_html=True)
- else:
- st.write("Result: ", generated_sequences[0], " (Wrong)", key="result_text", unsafe_allow_html=True)
-
-# Display correct grammar sentence in a box
-st.text("Correct Grammar Sentence:")
-st.code(generated_sequences[0])
+ st.write("### Results:")
+
+ # Check correctness and display in green or red
+ for generated_sequence in generated_sequences:
+ is_correct = generated_sequence == sent
+ color = "green" if is_correct else "red"
+ st.write(f"**Generated Sentence:**", generated_sequence, f" (Correct: {is_correct})", unsafe_allow_html=True)
+
+ # If incorrect, display correct grammar sentence in a box
+ if not is_correct:
+ st.warning(f"**Correct Grammar:** {sent}")
+
+# Display original input
+st.write("### Original Input:")
+st.write(sent)