Update app.py

app.py

@@ -11,6 +11,7 @@ import shutil
 import requests
 import pandas as pd
 import difflib
+from concurrent.futures import ThreadPoolExecutor
 
 # Define the device
 device = "cuda" if torch.cuda.is_available() else "cpu"
@@ -168,24 +169,32 @@ def split_text(text, max_tokens=500):
 
 
 # Function to generate text
-def ocr_correction(prompt, max_new_tokens=600):
-
+def ocr_correction(prompt, max_new_tokens=600, num_threads=os.cpu_count()):
     prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
+    # Set the number of threads for PyTorch
+    torch.set_num_threads(num_threads)
+
     # Generate text
-    output = model.generate(input_ids,
-                            max_new_tokens=max_new_tokens,
-                            pad_token_id=tokenizer.eos_token_id,
-                            top_k=50)
+    with ThreadPoolExecutor(max_workers=num_threads) as executor:
+        future = executor.submit(
+            model.generate,
+            input_ids,
+            max_new_tokens=max_new_tokens,
+            pad_token_id=tokenizer.eos_token_id,
+            top_k=50,
+            num_return_sequences=1,
+            do_sample=True,
+            temperature=0.7
+        )
+        output = future.result()
 
     # Decode and return the generated text
     result = tokenizer.decode(output[0], skip_special_tokens=True)
-
     print(result)
-
+    
     result = result.split("### Correction ###")[1]
-
     return result
 
 # OCR Correction Class