Update app.py

app.py

@@ -4,6 +4,7 @@ from transformers import GPT2LMHeadModel, GPT2Tokenizer
 import torch
 import gradio as gr
 import difflib
+import spaces
 from concurrent.futures import ThreadPoolExecutor
 import os
 
@@ -56,23 +57,18 @@ def split_text(text, max_tokens=400):
 
     return chunks
 
+@spaces.GPU
 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)
 
-    torch.set_num_threads(num_threads)
-
-    with ThreadPoolExecutor(max_workers=num_threads) as executor:
-        future = executor.submit(
-            model.generate,
-            input_ids,
+    output = 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=False
         )
-        output = future.result()
 
     result = tokenizer.decode(output[0], skip_special_tokens=True)
     return result.split("### Correction ###")[1].strip()