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)