file name cleanup (#941)

refuel-ai · Dec 4, 2024 · 25d51d4 · 25d51d4
1 parent 84f5762
commit 25d51d4
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/autolabel/transforms/ocr.py b/src/autolabel/transforms/ocr.py
@@ -19,7 +19,6 @@
 
 
 class OCRTransform(BaseTransform):
-
     """
     Extract text from documents using OCR.
 
@@ -126,7 +125,7 @@ def download_file(self, file_location: str) -> str:
 
         import requests
 
-        ext = Path(file_location).suffix
+        ext = Path(file_location).suffix.split("?")[0]
         with tempfile.NamedTemporaryFile(suffix=ext, delete=False) as temp_file:
             response = requests.get(file_location, timeout=30)
             response.raise_for_status()