Update app.py

app.py

@@ -5,41 +5,8 @@ from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, GenerationConfig
 model = AutoModelForSeq2SeqLM.from_pretrained("vennify/t5-base-grammar-correction")
 tokenizer = AutoTokenizer.from_pretrained("vennify/t5-base-grammar-correction")
 
-def correct_text(text, max_length, min_length, max_new_tokens, min_new_tokens, num_beams, temperature, top_p):
-    inputs = tokenizer.encode("grammar: " + text, return_tensors="pt")
-    
-    if max_new_tokens > 0 or min_new_tokens > 0:
-        if max_new_tokens > 0 and min_new_tokens > 0:
-            outputs = model.generate(
-                inputs,
-                max_new_tokens=max_new_tokens,
-                min_new_tokens=min_new_tokens,
-                num_beams=num_beams,
-                temperature=temperature,
-                top_p=top_p,
-                early_stopping=True,
-                do_sample=True
-            )
-        elif max_new_tokens > 0:
-            outputs = model.generate(inputs, max_new_tokens=max_new_tokens, min_length=min_length, num_beams=num_beams, temperature=temperature, top_p=top_p, early_stopping=True, do_sample=True)
-        else:
-            outputs = model.generate(inputs, max_length=max_length, min_new_tokens=min_new_tokens, num_beams=num_beams, temperature=temperature, top_p=top_p, early_stopping=True, do_sample=True)
-    else:
-        outputs = model.generate(
-            inputs,
-            max_length=max_length,
-            min_length=min_length,
-            num_beams=num_beams,
-            temperature=temperature,
-            top_p=top_p,
-            early_stopping=True,
-            do_sample=True
-        )
-    
-    corrected_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    yield corrected_text
 
-def correct_text2(text, genConfig):
+def correct_text(text, genConfig):
     inputs = tokenizer.encode("grammar: " + text, return_tensors="pt")
     outputs = model.generate(inputs, **genConfig.to_dict())
 
@@ -64,11 +31,9 @@ def respond(text, max_length, min_length, max_new_tokens, min_new_tokens, num_be
     if min_new_tokens > 0: 
         config.min_new_tokens = min_new_tokens
     
-    corrected = correct_text2(text, config)
+    corrected = correct_text(text, config)
     yield corrected
 
-    
-
 def update_prompt(prompt):
     return prompt
 
@@ -87,7 +52,6 @@ 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):
@@ -102,6 +66,6 @@ with gr.Blocks() as demo:
         
 
     
-    submitBtn.click(respond, [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_k, top_p], output_box)
 
 demo.launch()