Update scripts/document_tool.py

scripts/document_tool.py

@@ -13,6 +13,7 @@ import os
 import re
 import time
 import tempfile
+import spaces
 import numpy as np
 from tqdm import tqdm
 
@@ -400,6 +401,7 @@ class LegalDocumentTool(Tool):
             # Clean up the temporary file
             os.remove(temp_path)
 
+    @spaces.GPU
     def forward(
         self,
         text: Optional[str] = None,