refactor(seed): package unstructured seed reader as local Data Designer plugin
This commit is contained in:
parent
95dd202ab3
commit
c88cce8185
10 changed files with 115 additions and 96 deletions
3
setup.sh
3
setup.sh
|
|
@ -168,6 +168,7 @@ SINGLE_ENV_CONSTRAINTS="$REQ_ROOT/single-env/constraints.txt"
|
|||
SINGLE_ENV_DATA_DESIGNER="$REQ_ROOT/single-env/data-designer.txt"
|
||||
SINGLE_ENV_DATA_DESIGNER_DEPS="$REQ_ROOT/single-env/data-designer-deps.txt"
|
||||
SINGLE_ENV_PATCH="$REQ_ROOT/single-env/patch_metadata.py"
|
||||
LOCAL_DD_UNSTRUCTURED_PLUGIN="$SCRIPT_DIR/studio/backend/plugins/data-designer-unstructured-seed"
|
||||
|
||||
install_python_stack() {
|
||||
run_quiet "pip upgrade" pip install --upgrade pip
|
||||
|
|
@ -191,6 +192,8 @@ install_python_stack() {
|
|||
run_quiet "pip install data-designer deps" pip install --no-cache-dir -c "$SINGLE_ENV_CONSTRAINTS" -r "$SINGLE_ENV_DATA_DESIGNER_DEPS"
|
||||
echo " Installing data-designer..."
|
||||
run_quiet "pip install data-designer" pip install --no-cache-dir --no-deps -c "$SINGLE_ENV_CONSTRAINTS" -r "$SINGLE_ENV_DATA_DESIGNER"
|
||||
echo " Installing local data-designer unstructured plugin..."
|
||||
run_quiet "pip install data-designer-unstructured-seed" pip install --no-cache-dir --no-deps -e "$LOCAL_DD_UNSTRUCTURED_PLUGIN"
|
||||
run_quiet "patch single-env metadata" python "$SINGLE_ENV_PATCH"
|
||||
run_quiet "pip check" pip check
|
||||
echo "✅ Python dependencies installed"
|
||||
|
|
|
|||
|
|
@ -7,8 +7,6 @@ from pathlib import Path
|
|||
from typing import Any
|
||||
|
||||
from .jsonable import to_jsonable
|
||||
from .unstructured_seed_plugin import ensure_unstructured_seed_plugin_registered
|
||||
|
||||
_IMAGE_CONTEXT_PATCHED = False
|
||||
|
||||
|
||||
|
|
@ -100,8 +98,6 @@ def _apply_data_designer_image_context_patch() -> None:
|
|||
if _IMAGE_CONTEXT_PATCHED:
|
||||
return
|
||||
|
||||
ensure_unstructured_seed_plugin_registered()
|
||||
|
||||
try:
|
||||
from data_designer.config.models import ImageContext
|
||||
except ImportError:
|
||||
|
|
|
|||
|
|
@ -1,90 +0,0 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
|
||||
from .unstructured_seed import (
|
||||
DEFAULT_CHUNK_OVERLAP,
|
||||
DEFAULT_CHUNK_SIZE,
|
||||
materialize_unstructured_seed_dataset,
|
||||
resolve_chunking,
|
||||
)
|
||||
|
||||
try:
|
||||
import data_designer.lazy_heavy_imports as lazy
|
||||
from data_designer.config.seed_source import SeedSource
|
||||
from data_designer.engine.resources.seed_reader import SeedReader
|
||||
except ImportError: # pragma: no cover
|
||||
lazy = None
|
||||
|
||||
class SeedSource: # type: ignore[no-redef]
|
||||
pass
|
||||
|
||||
class SeedReader: # type: ignore[no-redef]
|
||||
@classmethod
|
||||
def __class_getitem__(cls, _item):
|
||||
return cls
|
||||
|
||||
|
||||
class UnstructuredSeedSource(SeedSource):
|
||||
seed_type: Literal["unstructured"] = "unstructured"
|
||||
path: str = Field(..., min_length=1)
|
||||
chunk_size: int = DEFAULT_CHUNK_SIZE
|
||||
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP
|
||||
|
||||
@field_validator("path", mode="after")
|
||||
@classmethod
|
||||
def _validate_path(cls, value: str) -> str:
|
||||
path = Path(value).expanduser()
|
||||
if not path.is_file():
|
||||
raise ValueError(f"Unstructured seed path is not a file: {path}")
|
||||
return value
|
||||
|
||||
@field_validator("chunk_size", mode="after")
|
||||
@classmethod
|
||||
def _validate_chunk_size(cls, value: int) -> int:
|
||||
size, _ = resolve_chunking(value, 0)
|
||||
return size
|
||||
|
||||
@field_validator("chunk_overlap", mode="after")
|
||||
@classmethod
|
||||
def _validate_chunk_overlap(cls, value: int, info) -> int:
|
||||
size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default)
|
||||
_, overlap = resolve_chunking(size, value)
|
||||
return overlap
|
||||
|
||||
|
||||
class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]):
|
||||
def create_duckdb_connection(self):
|
||||
if lazy is None:
|
||||
raise RuntimeError("data_designer is not available")
|
||||
return lazy.duckdb.connect()
|
||||
|
||||
def get_dataset_uri(self) -> str:
|
||||
path, _ = materialize_unstructured_seed_dataset(
|
||||
source_path=Path(self.source.path),
|
||||
chunk_size=self.source.chunk_size,
|
||||
chunk_overlap=self.source.chunk_overlap,
|
||||
)
|
||||
return str(path)
|
||||
|
||||
|
||||
def ensure_unstructured_seed_plugin_registered() -> None:
|
||||
try:
|
||||
from data_designer.plugins.plugin import Plugin, PluginType
|
||||
from data_designer.plugins.registry import PluginRegistry
|
||||
except ImportError:
|
||||
return
|
||||
|
||||
registry = PluginRegistry()
|
||||
if registry.plugin_exists("unstructured"):
|
||||
return
|
||||
|
||||
plugin = Plugin(
|
||||
impl_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedReader",
|
||||
config_qualified_name="core.data_recipe.unstructured_seed_plugin.UnstructuredSeedSource",
|
||||
plugin_type=PluginType.SEED_READER,
|
||||
)
|
||||
registry._plugins[plugin.name] = plugin # type: ignore[attr-defined]
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
[build-system]
|
||||
requires = ["setuptools>=68", "wheel"]
|
||||
build-backend = "setuptools.build_meta"
|
||||
|
||||
[project]
|
||||
name = "data-designer-unstructured-seed"
|
||||
version = "0.1.0"
|
||||
description = "Local Data Designer unstructured seed reader plugin"
|
||||
requires-python = ">=3.11"
|
||||
dependencies = [
|
||||
"data-designer-engine>=0.5.1,<0.6",
|
||||
"pandas>=2,<3",
|
||||
]
|
||||
|
||||
[project.entry-points."data_designer.plugins"]
|
||||
unstructured = "data_designer_unstructured_seed.plugin:unstructured_seed_plugin"
|
||||
|
||||
[tool.setuptools]
|
||||
package-dir = {"" = "src"}
|
||||
|
||||
[tool.setuptools.packages.find]
|
||||
where = ["src"]
|
||||
|
|
@ -0,0 +1,21 @@
|
|||
from .chunking import (
|
||||
DEFAULT_CHUNK_OVERLAP,
|
||||
DEFAULT_CHUNK_SIZE,
|
||||
build_unstructured_preview_rows,
|
||||
materialize_unstructured_seed_dataset,
|
||||
resolve_chunking,
|
||||
)
|
||||
from .config import UnstructuredSeedSource
|
||||
from .impl import UnstructuredSeedReader
|
||||
from .plugin import unstructured_seed_plugin
|
||||
|
||||
__all__ = [
|
||||
"DEFAULT_CHUNK_OVERLAP",
|
||||
"DEFAULT_CHUNK_SIZE",
|
||||
"build_unstructured_preview_rows",
|
||||
"materialize_unstructured_seed_dataset",
|
||||
"resolve_chunking",
|
||||
"UnstructuredSeedSource",
|
||||
"UnstructuredSeedReader",
|
||||
"unstructured_seed_plugin",
|
||||
]
|
||||
|
|
@ -5,7 +5,6 @@ import re
|
|||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
|
||||
DEFAULT_CHUNK_SIZE = 1200
|
||||
DEFAULT_CHUNK_OVERLAP = 200
|
||||
MAX_CHUNK_SIZE = 20000
|
||||
|
|
@ -44,6 +43,7 @@ def build_unstructured_preview_rows(
|
|||
import pandas as pd
|
||||
except ImportError as exc: # pragma: no cover
|
||||
raise RuntimeError(f"pandas is required for unstructured seed processing: {exc}") from exc
|
||||
|
||||
dataframe = pd.read_parquet(parquet_path).head(count)
|
||||
return [
|
||||
{"chunk_text": str(value.get("chunk_text", "")).strip()}
|
||||
|
|
@ -0,0 +1,38 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
from typing import Literal
|
||||
|
||||
from pydantic import Field, field_validator
|
||||
|
||||
from data_designer.config.seed_source import SeedSource
|
||||
|
||||
from .chunking import DEFAULT_CHUNK_OVERLAP, DEFAULT_CHUNK_SIZE, resolve_chunking
|
||||
|
||||
|
||||
class UnstructuredSeedSource(SeedSource):
|
||||
seed_type: Literal["unstructured"] = "unstructured"
|
||||
path: str = Field(..., min_length=1)
|
||||
chunk_size: int = DEFAULT_CHUNK_SIZE
|
||||
chunk_overlap: int = DEFAULT_CHUNK_OVERLAP
|
||||
|
||||
@field_validator("path", mode="after")
|
||||
@classmethod
|
||||
def _validate_path(cls, value: str) -> str:
|
||||
path = Path(value).expanduser()
|
||||
if not path.is_file():
|
||||
raise ValueError(f"Unstructured seed path is not a file: {path}")
|
||||
return value
|
||||
|
||||
@field_validator("chunk_size", mode="after")
|
||||
@classmethod
|
||||
def _validate_chunk_size(cls, value: int) -> int:
|
||||
size, _ = resolve_chunking(value, 0)
|
||||
return size
|
||||
|
||||
@field_validator("chunk_overlap", mode="after")
|
||||
@classmethod
|
||||
def _validate_chunk_overlap(cls, value: int, info) -> int:
|
||||
size = info.data.get("chunk_size", cls.model_fields["chunk_size"].default)
|
||||
_, overlap = resolve_chunking(size, value)
|
||||
return overlap
|
||||
|
|
@ -0,0 +1,22 @@
|
|||
from __future__ import annotations
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import data_designer.lazy_heavy_imports as lazy
|
||||
from data_designer.engine.resources.seed_reader import SeedReader
|
||||
|
||||
from .chunking import materialize_unstructured_seed_dataset
|
||||
from .config import UnstructuredSeedSource
|
||||
|
||||
|
||||
class UnstructuredSeedReader(SeedReader[UnstructuredSeedSource]):
|
||||
def create_duckdb_connection(self):
|
||||
return lazy.duckdb.connect()
|
||||
|
||||
def get_dataset_uri(self) -> str:
|
||||
path, _ = materialize_unstructured_seed_dataset(
|
||||
source_path=Path(self.source.path),
|
||||
chunk_size=self.source.chunk_size,
|
||||
chunk_overlap=self.source.chunk_overlap,
|
||||
)
|
||||
return str(path)
|
||||
|
|
@ -0,0 +1,7 @@
|
|||
from data_designer.plugins.plugin import Plugin, PluginType
|
||||
|
||||
unstructured_seed_plugin = Plugin(
|
||||
impl_qualified_name="data_designer_unstructured_seed.impl.UnstructuredSeedReader",
|
||||
config_qualified_name="data_designer_unstructured_seed.config.UnstructuredSeedSource",
|
||||
plugin_type=PluginType.SEED_READER,
|
||||
)
|
||||
|
|
@ -10,7 +10,7 @@ from typing import Any
|
|||
from uuid import uuid4
|
||||
|
||||
from fastapi import APIRouter, HTTPException
|
||||
from core.data_recipe.unstructured_seed import (
|
||||
from data_designer_unstructured_seed.chunking import (
|
||||
build_unstructured_preview_rows,
|
||||
resolve_chunking,
|
||||
)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue