Update app.py

app.py

@@ -1,6 +1,12 @@
 # app.py
 from transformers import T5Tokenizer, T5ForConditionalGeneration
 import gradio as gr
+import nltk
+from nltk.tokenize import sent_tokenize
+import difflib
+
+# Download the punkt tokenizer for sentence splitting
+nltk.download('punkt')
 
 # Load a pre-trained T5 model specifically fine-tuned for grammar correction
 tokenizer = T5Tokenizer.from_pretrained("prithivida/grammar_error_correcter_v1")
@@ -8,28 +14,53 @@ model = T5ForConditionalGeneration.from_pretrained("prithivida/grammar_error_cor
 
 # Function to perform grammar correction
 def grammar_check(text):
-    input_text = f"gec: {text}"
-    input_ids = tokenizer.encode(input_text, return_tensors="pt")
-    outputs = model.generate(input_ids)
-    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # Split the text into sentences
+    sentences = sent_tokenize(text)
+    corrected_sentences = []
+    original_sentences = []
+
+    for sentence in sentences:
+        original_sentences.append(sentence)
+        input_text = f"gec: {sentence}"
+        input_ids = tokenizer.encode(input_text, return_tensors="pt", max_length=512, truncation=True)
+        outputs = model.generate(input_ids, max_length=512, num_beams=4, early_stopping=True)
+        corrected_sentence = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        corrected_sentences.append(corrected_sentence)
+
+    # Function to underline and color revised parts
+    def underline_and_color_revisions(original, corrected):
+        diff = difflib.ndiff(original.split(), corrected.split())
+        result = []
+        for word in diff:
+            if word.startswith("+ "):
+                result.append(f"<u style='color:red;'>{word[2:]}</u>")
+            elif word.startswith("- "):
+                continue
+            else:
+                result.append(word[2:])
+        return " ".join(result)
+
+    # Join the corrected sentences back into a single string
+    corrected_text = " ".join(
+        underline_and_color_revisions(orig, corr) for orig, corr in zip(original_sentences, corrected_sentences)
+    )
     return corrected_text
 
 # Create Gradio interface with a writing prompt
 interface = gr.Interface(
     fn=grammar_check,
     inputs="text",
-    outputs="text",
+    outputs="html",  # Change output type to HTML
     title="Grammar Checker",
     description=(
         "Enter text to check for grammar mistakes.\n\n"
         "Writing Prompt:\n"
-        "In the story, Alex and his friends decided to donate all the treasure to the town library. "
-        "This can be a good metaphor for the idea that individuals should sacrifice for society. Some people "
-        "claim that individuals do not have to sacrifice but only think about their own happiness. To what extent "
-        "do you agree or disagree with this?\n\n"
-        "Explain your reasons with some relevant examples from your own experience and knowledge. "
+        "In the story, Alex and his friends discovered an ancient treasure in Whispering Hollow and decided to donate the artifacts to the local museum.\n\n"
+        "In the past, did you have a similar experience where you found something valuable or interesting? Tell the story. Describe what you found, what you did with it, and how you felt about your decision.\n\n"
+        "Remember to use past tense in your writing.\n"
+        "Sample text for testing: When I was 10, I find an old coin in my backyard. I kept it for a while and shows it to my friends. They was impressed and say it might be valuable. Later, I take it to a local antique shop, and the owner told me it was very old. I decided to give it to the museum in my town. The museum was happy and put it on display. I feel proud of my decision."
     )
 )
 
 # Launch the interface
-interface.launch()
+interface.launch()