Update app.py

app.py

@@ -1,5 +1,5 @@
 import gradio as gr
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig
 
 # Load the model and tokenizer
 model = AutoModelForSeq2SeqLM.from_pretrained("vennify/t5-base-grammar-correction")
@@ -39,19 +39,43 @@ def correct_text(text, max_length, min_length, max_new_tokens, min_new_tokens, n
     corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
     yield corrected_text
 
+def correct_text2(text, genConfig):
+    inputs = tokenizer.encode("grammar: " + text, return_tensors="pt")
+    outputs = model.generate(inputs, **genConfig.to_dict())
+
+    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    yield corrected_text
+
+def respond(text, max_length, min_length, max_new_tokens, min_new_tokens, num_beams, temperature, top_p):
+    config = GenerationConfig(
+        max_length=max_length,
+        min_length=min_length,
+        num_beams=num_beams,
+        temperature=temperature,
+        top_p=top_p,
+        early_stopping=True,
+        do_sample=True
+    )
+
+    # Add max/min new tokens if they are there
+    if max_new_tokens > 0: 
+        config.max_new_tokens = max_new_tokens
+    if min_new_tokens > 0: 
+        config.min_new_tokens = min_new_tokens
+    
+    corrected = correct_text2(text, config)
+    yield corrected
+
+    
 
 def update_prompt(prompt):
     return prompt
 
 # Create the Gradio interface
 with gr.Blocks() as demo:
-    gr.Markdown(
-    """
-    # Grammar Correction App
-    """)
-    prompt_box = gr.Textbox(lines=2, placeholder="Enter your prompt here...")
+    gr.Markdown("""# Grammar Correction App""")
+    prompt_box = gr.Textbox(placeholder="Enter your prompt here...")
     output_box = gr.Textbox()
-    submitBtn = gr.Button("Submit")
 
     # Sample prompts
     with gr.Row():
@@ -62,7 +86,8 @@ with gr.Blocks() as demo:
         samp1.click(update_prompt, samp1, prompt_box)
         samp2.click(update_prompt, samp2, prompt_box)
         samp3.click(update_prompt, samp3, prompt_box)
-        
+
+    submitBtn = gr.Button("Submit")
     
     with gr.Accordion("Generation Parameters:", open=False):
         max_length = gr.Slider(minimum=1, maximum=256, value=80, step=1, label="Max Length")
@@ -74,8 +99,6 @@ with gr.Blocks() as demo:
         top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (nucleus sampling)")
 
     
-    submitBtn.click(correct_text, [prompt_box, max_length, min_length, max_tokens, min_tokens, num_beams, temperature, top_p], output_box)
-
-    
+    submitBtn.click(respond, [prompt_box, max_length, min_length, max_tokens, min_tokens, num_beams, temperature, top_p], output_box)
 
-demo.launch()
+demo.launch()