Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
31 changes: 31 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
name: Tests

on:
push:
pull_request:
workflow_dispatch:

jobs:
pytest:
name: Pytest (${{ matrix.python-version }})
runs-on: ubuntu-latest
strategy:
matrix:
python-version: ["3.10", "3.11", "3.12", "3.13"]

steps:
- uses: actions/checkout@v4

- name: Set up Python
uses: actions/setup-python@v5
with:
python-version: ${{ matrix.python-version }}
cache: "pip"

- name: Install dependencies
run: |
python -m pip install --upgrade pip
python -m pip install -e . pytest

- name: Run tests
run: python -m pytest tests
11 changes: 11 additions & 0 deletions deepdoc/common/misc_utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@
#

import logging
import os

logger = logging.getLogger(__name__)

Expand All @@ -31,3 +32,13 @@ def pip_install_torch():
return True
except ImportError:
return False


def parse_bool(value: str | None, default: bool = False) -> bool:
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}


def offline_mode_or_from_env(offline: bool | None = None) -> bool:
return offline if offline is not None else parse_bool(os.getenv("DEEPDOC_OFFLINE"), default=False)
27 changes: 18 additions & 9 deletions deepdoc/common/model_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,17 +22,14 @@
import logging
import os
from dataclasses import dataclass
from importlib import resources
from pathlib import Path
from ..common.misc_utils import offline_mode_or_from_env


GLOBAL_MODELSCOPE_REPO_ENV = "DEEPDOC_MODELSCOPE_REPO"
GLOBAL_MODELSCOPE_REVISION_ENV = "DEEPDOC_MODELSCOPE_REVISION"


def _parse_bool(value: str | None, default: bool = False) -> bool:
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}
TOKENIZER_MODEL_DIR_ENV = "DEEPDOC_TOKENIZER_MODEL_DIR"


def _normalize_provider(provider: str | None) -> str:
Expand All @@ -56,6 +53,18 @@ def _model_home_path(model_home: str | None) -> Path:
return Path.home().joinpath(".cache", "deepdoc")


def _resolve_tokenizer_dict_path() -> Path:
configured_dir = os.getenv(TOKENIZER_MODEL_DIR_ENV)
if configured_dir:
dictionary = Path(configured_dir).expanduser().resolve().joinpath("huqie.txt")
else:
dictionary = Path(str(resources.files("deepdoc").joinpath("dict", "huqie.txt"))).resolve()

if not dictionary.exists():
raise FileNotFoundError("Tokenizer dictionary not found: {}. Set {} to a directory containing huqie.txt.".format(dictionary, TOKENIZER_MODEL_DIR_ENV))
return dictionary


@dataclass(frozen=True)
class BundleSpec:
name: str
Expand Down Expand Up @@ -243,7 +252,7 @@ def resolve_bundle_dir(

spec = BUNDLES[bundle]
provider_name = _normalize_provider(provider)
offline_mode = offline if offline is not None else _parse_bool(os.getenv("DEEPDOC_OFFLINE"), default=False)
offline_mode = offline_mode_or_from_env(offline)

explicit_local = os.getenv(spec.local_dir_env)
if explicit_local:
Expand Down Expand Up @@ -356,5 +365,5 @@ def resolve_tokenizer_dict_prefix(
provider: str | None = None,
offline: bool | None = None,
) -> str:
bundle_dir = Path(resolve_bundle_dir("tokenizer", model_home=model_home, provider=provider, offline=offline))
return str(bundle_dir.joinpath("huqie"))
del model_home, provider, offline
return str(_resolve_tokenizer_dict_path().with_suffix(""))
25 changes: 6 additions & 19 deletions deepdoc/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,16 +11,11 @@
resolve_xgb_model_dir,
validate_bundle_dir,
)
from .common.misc_utils import offline_mode_or_from_env

ProviderType = Literal["local", "modelscope", "auto"]


def _parse_bool(value: str | None, default: bool = False) -> bool:
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}


def _normalize_provider(provider: str) -> ProviderType:
normalized = provider.strip().lower()
aliases = {
Expand All @@ -31,9 +26,7 @@ def _normalize_provider(provider: str) -> ProviderType:
}
normalized = aliases.get(normalized, normalized)
if normalized not in {"local", "modelscope", "auto"}:
raise ValueError(
"Unsupported model provider '{}'. Use one of: local, modelscope, auto.".format(provider)
)
raise ValueError("Unsupported model provider '{}'. Use one of: local, modelscope, auto.".format(provider))
return normalized # type: ignore[return-value]


Expand Down Expand Up @@ -64,13 +57,10 @@ def resolve_dict_path(self) -> str:
if dictionary.is_dir():
dictionary = dictionary.joinpath("huqie.txt")
if dictionary.suffix != ".txt":
raise ValueError(
"TokenizerConfig.dict_path must point to a '.txt' dictionary file, got: {}".format(dictionary)
)
raise ValueError("TokenizerConfig.dict_path must point to a '.txt' dictionary file, got: {}".format(dictionary))
_require_file(
dictionary,
"Tokenizer dictionary not found: {}. Provide a valid TokenizerConfig.dict_path."
.format(dictionary),
"Tokenizer dictionary not found: {}. Provide a valid TokenizerConfig.dict_path.".format(dictionary),
)
return str(dictionary)

Expand All @@ -86,7 +76,7 @@ def from_env(cls) -> "TokenizerConfig":

return cls(
dict_path=dict_path,
offline=_parse_bool(os.getenv("DEEPDOC_OFFLINE"), default=False),
offline=offline_mode_or_from_env(None),
nltk_data_dir=os.getenv("DEEPDOC_NLTK_DATA_DIR"),
)

Expand All @@ -107,10 +97,7 @@ def _resolve_bundle_dir(self, bundle: str, explicit_dir: str | None) -> str:
candidate = Path(explicit_dir).expanduser().resolve()
exists, missing = validate_bundle_dir(bundle, candidate)
if not exists:
raise FileNotFoundError(
"Missing required files for '{}' bundle in {}: {}"
.format(bundle, candidate, ", ".join(missing))
)
raise FileNotFoundError("Missing required files for '{}' bundle in {}: {}".format(bundle, candidate, ", ".join(missing)))
return str(candidate)

model_provider = self.normalized_provider()
Expand Down
25 changes: 5 additions & 20 deletions deepdoc/depend/nltk_manager.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,19 +2,13 @@
import os
import threading
from pathlib import Path
from ..common.misc_utils import offline_mode_or_from_env

logger = logging.getLogger(__name__)

_RESOURCE_SPECS: tuple[tuple[str, tuple[str, ...]], ...] = (
(
"punkt",
(
"tokenizers/punkt",
"tokenizers/punkt.zip",
"tokenizers/punkt_tab",
"tokenizers/punkt_tab.zip",
),
),
("punkt", ("tokenizers/punkt", "tokenizers/punkt.zip")),
("punkt_tab", ("tokenizers/punkt_tab", "tokenizers/punkt_tab.zip")),
("wordnet", ("corpora/wordnet", "corpora/wordnet.zip")),
(
"averaged_perceptron_tagger",
Expand All @@ -31,12 +25,6 @@
_ensured_keys: set[tuple[str, bool]] = set()


def _parse_bool(value: str | None, default: bool = False) -> bool:
if value is None:
return default
return value.strip().lower() in {"1", "true", "yes", "on"}


def _resolve_nltk_data_dir(data_dir: str | None) -> Path | None:
# Resolution precedence:
# 1) explicit arg
Expand Down Expand Up @@ -88,7 +76,7 @@ def ensure_nltk_data(
import nltk

resolved_dir = _resolve_nltk_data_dir(data_dir)
offline_mode = offline if offline is not None else _parse_bool(os.getenv("DEEPDOC_OFFLINE"), default=False)
offline_mode = offline_mode_or_from_env(offline)
auto_download_mode = not offline_mode

_ensure_search_path(nltk, resolved_dir)
Expand Down Expand Up @@ -117,10 +105,7 @@ def ensure_nltk_data(
if missing_packages:
searched_paths = ", ".join(nltk.data.path)
raise RuntimeError(
"Missing required NLTK packages: {}. Searched paths: {}. "
"Set DEEPDOC_NLTK_DATA_DIR to a local NLTK data path, or disable offline mode by setting "
"DEEPDOC_OFFLINE=0."
.format(
"Missing required NLTK packages: {}. Searched paths: {}. Set DEEPDOC_NLTK_DATA_DIR to a local NLTK data path, or disable offline mode by setting DEEPDOC_OFFLINE=0.".format(
", ".join(missing_packages),
searched_paths,
)
Expand Down
5 changes: 5 additions & 0 deletions pyproject.toml
Original file line number Diff line number Diff line change
Expand Up @@ -89,6 +89,11 @@ extend-select = ["ASYNC", "ASYNC1"]
ignore = ["E402"]

[tool.pytest.ini_options]
addopts = [
"--ignore=tests/test_data",
"--ignore=tests/test_results",
]
testpaths = ["tests"]
markers = [
"p1: high priority test cases",
"p2: medium priority test cases",
Expand Down
24 changes: 16 additions & 8 deletions tests/test_model_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,10 +22,6 @@ def _create_combined_repo_layout(root: Path) -> None:
for name in ms.BUNDLES["xgb"].required_files:
_touch(root / "xgb" / name)

# Tokenizer bundle
for name in ms.BUNDLES["tokenizer"].required_files:
_touch(root / "tokenizer" / name)


class TestModelStoreSharedRepo(unittest.TestCase):
def setUp(self) -> None:
Expand Down Expand Up @@ -68,12 +64,10 @@ def snapshot_download(
with patch.object(ms, "_import_modelscope_snapshot_download", return_value=snapshot_download):
vision_dir = Path(ms.resolve_bundle_dir("vision", model_home=tmp, provider="modelscope", offline=False))
xgb_dir = Path(ms.resolve_bundle_dir("xgb", model_home=tmp, provider="modelscope", offline=False))
tok_dir = Path(ms.resolve_bundle_dir("tokenizer", model_home=tmp, provider="modelscope", offline=False))

expected_root = (Path(tmp) / "modelscope" / "Xorbits__deepdoc" / "v1").resolve()
self.assertEqual(vision_dir.resolve(), (expected_root / "vision").resolve())
self.assertEqual(xgb_dir.resolve(), (expected_root / "xgb").resolve())
self.assertEqual(tok_dir.resolve(), (expected_root / "tokenizer").resolve())

self.assertGreaterEqual(len(calls), 1)
for call in calls:
Expand Down Expand Up @@ -134,8 +128,22 @@ def snapshot_download(*args, **kwargs) -> str: # pragma: no cover
with patch.object(ms, "_import_modelscope_snapshot_download", return_value=snapshot_download):
vision_dir = Path(ms.resolve_bundle_dir("vision", model_home=tmp, provider="auto", offline=False))
xgb_dir = Path(ms.resolve_bundle_dir("xgb", model_home=tmp, provider="auto", offline=False))
tok_dir = Path(ms.resolve_bundle_dir("tokenizer", model_home=tmp, provider="auto", offline=False))

self.assertEqual(vision_dir.resolve(), (expected_root / "vision").resolve())
self.assertEqual(xgb_dir.resolve(), (expected_root / "xgb").resolve())
self.assertEqual(tok_dir.resolve(), (expected_root / "tokenizer").resolve())

def test_resolve_tokenizer_dict_prefix_uses_packaged_dict_by_default(self) -> None:
prefix = Path(ms.resolve_tokenizer_dict_prefix())

self.assertEqual(prefix.name, "huqie")
self.assertTrue(prefix.with_suffix(".txt").exists())

def test_resolve_tokenizer_dict_prefix_uses_env_dir_when_set(self) -> None:
with tempfile.TemporaryDirectory() as tmp:
tokenizer_dir = Path(tmp) / "tokenizer"
_touch(tokenizer_dir / "huqie.txt")
os.environ[ms.TOKENIZER_MODEL_DIR_ENV] = str(tokenizer_dir)

prefix = Path(ms.resolve_tokenizer_dict_prefix())

self.assertEqual(prefix.resolve(), (tokenizer_dir / "huqie").resolve())
Loading
Loading