feat: support revision (#75)

Author: Moskize91

doc_page_extractor/extractor.py

@@ -36,8 +36,8 @@ class _PageExtractorImpls:
     def __init__(self, model: DeepSeekOCRModel) -> None:
         self._model: DeepSeekOCRModel = model
 
-    def download_models(self) -> None:
-        self._model.download()
+    def download_models(self, revision: str | None = None) -> None:
+        self._model.download(revision)
 
     def load_models(self) -> None:
         self._model.load()

doc_page_extractor/model.py

@@ -80,11 +80,13 @@ def __init__(
             if next_model_index == 0:
                 raise ValueError("No devices are enabled for model loading.")
 
-    def download(self) -> None:
+    def download(self, revision: str | None) -> None:
         with self._rwlock.gen_wlock():
             snapshot_download(
                 repo_id=self._model_name,
                 repo_type="model",
+                revision=revision,
+                force_download=True,
                 cache_dir=self._cache_dir(),
             )
             if self._model_path is not None and self._find_pretrained_path() is None:

doc_page_extractor/types.py

@@ -28,7 +28,7 @@ class ExtractionContext:
 
 @runtime_checkable
 class PageExtractor(Protocol):
-    def download_models(self) -> None:
+    def download_models(self, revision: str | None = None) -> None:
         ...
 
     def load_models(self) -> None:
@@ -47,7 +47,7 @@ def extract(
 
 @runtime_checkable
 class DeepSeekOCRModel(Protocol):
-    def download(self) -> None:
+    def download(self, revision: str | None) -> None:
         ...
 
     def load(self) -> None:

download.py

@@ -0,0 +1,20 @@
+import time
+from pathlib import Path
+
+from doc_page_extractor import create_page_extractor
+
+
+_REVISION = "9f30c71f441d010e5429c532364a86705536c53a"
+
+def main() -> None:
+    project_root = Path(__file__).parent
+    extractor = create_page_extractor(
+        model_path=project_root / "models-cache",
+        local_only=False,
+    )
+    begin_at = time.time()
+    extractor.download_models(_REVISION)
+    print(f"Models downloaded cost {time.time() - begin_at:.2f} seconds.")
+
+if __name__ == "__main__":
+    main()

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [project]
 name = "doc-page-extractor"
-version = "1.0.7"
+version = "1.0.8"
 description = "Document page extraction tool powered by DeepSeek-OCR"
 authors = [
     {name = "Tao Zeyu", email = "i@taozeyu.com"}