Update app.py

app.py

@@ -129,19 +129,25 @@ def ocr_correction(prompt, max_new_tokens=600):
     splits = split_text(prompt, max_tokens=400)
     corrected_splits = []
 
+    list_prompts = []
+
     for split in splits:
         full_prompt = f"### Text ###\n{split}\n\n\n### Correction ###\n"
+        print(full_prompt)
         encoded = tokenizer.encode(full_prompt)
         prompt_tokens = tokenizer.convert_ids_to_tokens(encoded)
-
-        result = generator.generate_batch(
-            [prompt_tokens],
-            max_length=max_new_tokens,
-            sampling_temperature=0.7,
-            sampling_topk=20,
-            include_prompt_in_result=False
-        )[0]
-
+        list_prompts.append(prompt_tokens)
+
+    results = generator.generate_batch(
+        list_prompts,
+        max_length=max_new_tokens,
+        sampling_temperature=0.7,
+        sampling_topk=20,
+        repetition_penalty=1.1,
+        include_prompt_in_result=False
+    )
+
+    for result in results:
         corrected_text = tokenizer.decode(result.sequences_ids[0])
         corrected_splits.append(corrected_text)