Update app.py

app.py

@@ -13,22 +13,19 @@ import pandas as pd
 import difflib
 from concurrent.futures import ThreadPoolExecutor
 
-# Define the device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
 # OCR Correction Model
 ocr_model_name = "PleIAs/OCRonos-Vintage"
 
 import torch
 from transformers import GPT2LMHeadModel, GPT2Tokenizer
 
+device = "cuda"
+
 # Load pre-trained model and tokenizer
 model_name = "PleIAs/OCRonos-Vintage"
 model = GPT2LMHeadModel.from_pretrained(model_name)
 tokenizer = GPT2Tokenizer.from_pretrained(model_name)
 
-# Set the device to GPU if available, otherwise use CPU
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 # CSS for formatting
@@ -169,7 +166,9 @@ def split_text(text, max_tokens=500):
 
 
 # Function to generate text
-def ocr_correction(prompt, max_new_tokens=600, num_threads=os.cpu_count()):
+@spaces.GPU
+def ocr_correction(prompt, max_new_tokens=500):
+    
     prompt = f"""### Text ###\n{prompt}\n\n\n### Correction ###\n"""
     input_ids = tokenizer.encode(prompt, return_tensors="pt").to(device)
 
@@ -177,9 +176,7 @@ def ocr_correction(prompt, max_new_tokens=600, num_threads=os.cpu_count()):
     torch.set_num_threads(num_threads)
 
     # Generate text
-    with ThreadPoolExecutor(max_workers=num_threads) as executor:
-        future = executor.submit(
-            model.generate,
+    output = model.generate,
             input_ids,
             max_new_tokens=max_new_tokens,
             pad_token_id=tokenizer.eos_token_id,
@@ -188,8 +185,6 @@ def ocr_correction(prompt, max_new_tokens=600, num_threads=os.cpu_count()):
             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)