amsdal_ml 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

amsdal_ml/ml_ingesting/loaders/pdf_loader.py

@@ -0,0 +1,136 @@
+from __future__ import annotations
+
+import asyncio
+import logging
+from typing import IO
+from typing import Any
+
+import pymupdf  # type: ignore
+
+from amsdal_ml.ml_ingesting.loaders.loader import Loader
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+from amsdal_ml.ml_ingesting.types import LoadedPage
+
+
+def _merge_spaced_characters(text: str) -> str:
+    """Collapse sequences like "s h o r t - t e r m" into "short-term".
+
+    Some PDFs return every character as a separate token. We merge contiguous
+    single-character tokens (letters/digits and simple inline punctuation) so
+    downstream cleaning does not preserve the artificial spaces.
+    """
+
+    tokens = text.replace("\n", " ").split()
+    merged: list[str] = []
+    buffer = ""
+
+    def flush_buffer() -> None:
+        nonlocal buffer
+        if buffer:
+            merged.append(buffer)
+            buffer = ""
+
+    for tok in tokens:
+        if len(tok) == 1 and (tok.isalnum() or tok in {"'", "-", "&"}):
+            buffer += tok
+            continue
+        flush_buffer()
+        merged.append(tok)
+
+    flush_buffer()
+    return " ".join(merged)
+
+
+def _is_noise(
+    text: str,
+    *,
+    min_tokens: int = 20,
+    single_char_ratio: float = 0.6,
+    single_char_min_count: int = 5,
+) -> bool:
+    """Heuristic: drop text dominated by single-character tokens.
+
+    We treat as noise when:
+    - Most tokens are single characters (ratio >= single_char_ratio), AND
+      either the sample is long enough (>= min_tokens) or we have at least a
+      few single-character tokens (>= single_char_min_count) even in short text.
+    """
+
+    tokens = text.split()
+    if not tokens:
+        return True
+
+    single_chars = sum(1 for t in tokens if len(t) == 1)
+    ratio = single_chars / len(tokens)
+
+    if ratio < single_char_ratio:
+        return False
+
+    return len(tokens) >= min_tokens or single_chars >= single_char_min_count
+
+
+class PdfLoader(Loader):
+    def __init__(self, *, extract_metadata: bool = True) -> None:
+        self.extract_metadata = extract_metadata
+        self.logger = logging.getLogger(__name__)
+
+    def _read_sync(
+        self,
+        file: IO[Any],
+        filename: str | None = None,
+        metadata: dict[str, Any] | None = None
+    ) -> LoadedDocument:
+        peek = file.read(4)
+        file.seek(0)
+        if not (isinstance(peek, bytes) and peek.startswith(b"%PDF")):
+            msg = "Not a PDF file"
+            raise ValueError(msg)
+
+        data = file.read()
+
+        with pymupdf.open(stream=data, filetype="pdf") as doc:
+            pages: list[LoadedPage] = []
+            for idx, page in enumerate(doc):
+                text_raw = page.get_text("text") or ''
+                text = _merge_spaced_characters(text_raw)
+                if not text:
+                    continue
+                if _is_noise(text):
+                    self.logger.debug("Skipping page %s as noise (single-char heavy)", idx + 1)
+                    continue
+                pages.append(
+                    LoadedPage(
+                        page_number=idx + 1,
+                        text=text,
+                        metadata={'page_number': idx + 1},
+                    )
+                )
+
+            doc_meta: dict[str, Any] = {}
+            if self.extract_metadata:
+                raw_meta = doc.metadata or {}
+                for key, value in raw_meta.items():
+                    if value is None:
+                        continue
+                    doc_meta[str(key)] = str(value)
+
+            if metadata:
+                doc_meta.update(metadata)
+            if filename:
+                doc_meta.setdefault('filename', filename)
+
+            return LoadedDocument(pages=pages, metadata=doc_meta)
+
+    def load(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        metadata: dict[str, Any] | None = None
+    ) -> LoadedDocument:
+        return self._read_sync(file, filename=filename, metadata=metadata)
+
+    async def aload(
+        self, file: IO[Any], *, filename: str | None = None, metadata: dict[str, Any] | None = None
+    ) -> LoadedDocument:
+        return await asyncio.to_thread(self._read_sync, file, filename, metadata)

amsdal_ml/ml_ingesting/loaders/text_loader.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+import asyncio
+from typing import IO
+from typing import Any
+
+from amsdal_ml.ml_ingesting.loaders.loader import Loader
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+from amsdal_ml.ml_ingesting.types import LoadedPage
+
+
+class TextLoader(Loader):
+    def __init__(self, *, encoding: str = 'utf-8', errors: str = 'ignore') -> None:
+        self.encoding = encoding
+        self.errors = errors
+
+    def _read_text(self, file: IO[Any]) -> str:
+        data = file.read()
+        if isinstance(data, bytes):
+            return data.decode(self.encoding, errors=self.errors)
+        return str(data)
+
+    def load(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> LoadedDocument:
+        text = self._read_text(file)
+        doc_meta = dict(metadata or {})
+        if filename:
+            doc_meta.setdefault('filename', filename)
+        page = LoadedPage(page_number=None, text=text, metadata={})
+        return LoadedDocument(pages=[page], metadata=doc_meta)
+
+    async def aload(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        metadata: dict[str, Any] | None = None,
+    ) -> LoadedDocument:
+        return await asyncio.to_thread(self.load, file, filename=filename, metadata=metadata)

amsdal_ml/ml_ingesting/model_ingester.py

@@ -0,0 +1,278 @@
+from __future__ import annotations
+
+import io
+import logging
+from collections.abc import AsyncIterator
+from collections.abc import Iterable
+from typing import IO
+from typing import Any
+from typing import ClassVar
+from typing import Generic
+from typing import TypeVar
+
+from amsdal.models.core.file import File
+from amsdal_models.querysets.base_queryset import QuerySetBase
+
+from amsdal_ml.ml_ingesting.loaders.loader import Loader
+from amsdal_ml.ml_ingesting.pipeline import DefaultIngestionPipeline
+from amsdal_ml.ml_ingesting.types import IngestionSource
+
+LoaderT = TypeVar('LoaderT', bound=Loader)
+
+
+class ModelIngester(Generic[LoaderT]):
+    MIN_OBJECTS_FOR_WARNING: ClassVar[int] = 3
+
+    def __init__(
+        self,
+        *,
+        pipeline: DefaultIngestionPipeline[LoaderT],
+        base_tags: Iterable[str] | None = None,
+        base_metadata: dict[str, Any] | None = None,
+        logger: logging.Logger | None = None,
+    ) -> None:
+        self.pipeline = pipeline
+        self.base_tags = list(base_tags or [])
+        self.base_metadata = dict(base_metadata or {})
+        self.logger = logger or logging.getLogger(__name__)
+
+    async def aingest(
+        self,
+        objects: Iterable[Any] | AsyncIterator[Any],
+        *,
+        fields: Iterable[str] | None = None,
+        tags: Iterable[str] | None = None,
+    ) -> list[Any]:
+        results: list[Any] = []
+        count_objects = 0
+        fields_set = set(fields) if fields else None
+
+        self.logger.debug("Starting async ingest; fields=%s", fields_set)
+
+        async for instance in self._amaterialize(objects):
+            count_objects += 1
+            object_class, object_id = self._resolve_link(instance)
+            self.logger.debug(
+                "Async ingest instance %s:%s; model_fields=%s",
+                object_class,
+                object_id,
+                getattr(instance.__class__, "model_fields", None),
+            )
+
+            found_any = False
+            async for payload in self._aiter_values(instance, fields_set):
+                found_any = True
+                field_name, field_title, stream, filename, kind = payload
+                source = self._build_source(object_class, object_id, field_name, field_title, filename)
+                run_tags = self._build_tags(tags, field_name, kind)
+                try:
+                    chunk_res = await self.pipeline.arun(
+                        stream,
+                        filename=filename,
+                        tags=run_tags,
+                        source=source,
+                    )
+                    results.extend(chunk_res)
+                except Exception as exc:  # noqa: BLE001
+                    self._warn_skip(object_class, field_name, exc)
+
+            if not found_any and count_objects <= self.MIN_OBJECTS_FOR_WARNING:
+                self.logger.info("  -> No ingestible fields found in %s:%s", object_class, object_id)
+
+        if count_objects == 0:
+            self.logger.warning("ModelIngester received 0 objects to process!")
+
+        self.logger.debug("Async ingest finished; processed=%s results=%s", count_objects, len(results))
+        return results
+
+    def ingest(
+        self,
+        objects: Iterable[Any],
+        *,
+        fields: Iterable[str] | None = None,
+        tags: Iterable[str] | None = None,
+    ) -> list[Any]:
+        rows = list(self._materialize(objects))
+        fields_set = set(fields) if fields else None
+        results: list[Any] = []
+        self.logger.debug("Starting sync ingest; rows=%s fields=%s", len(rows), fields_set)
+        for instance in rows:
+            object_class, object_id = self._resolve_link(instance)
+            self.logger.debug(
+                "Sync ingest instance %s:%s; model_fields=%s",
+                object_class,
+                object_id,
+                getattr(instance.__class__, "model_fields", None),
+            )
+            found_any = False
+            for payload in self._iter_values(instance, fields_set):
+                found_any = True
+                field_name, field_title, stream, filename, kind = payload
+                source = self._build_source(object_class, object_id, field_name, field_title, filename)
+                run_tags = self._build_tags(tags, field_name, kind)
+                try:
+                    results.extend(
+                        self.pipeline.run(
+                            stream,
+                            filename=filename,
+                            tags=run_tags,
+                            source=source,
+                        )
+                    )
+                except Exception as exc:  # noqa: BLE001
+                    self._warn_skip(object_class, field_name, exc)
+            if not found_any:
+                self.logger.debug("No ingestible fields in %s:%s", object_class, object_id)
+
+        self.logger.debug("Sync ingest finished; processed=%s results=%s", len(rows), len(results))
+        return results
+
+    async def _aiter_values(
+        self,
+        instance: Any,
+        fields_set: set[str] | None,
+    ) -> AsyncIterator[tuple[str, str | None, IO[Any], str | None, str]]:
+        if isinstance(instance, File):
+            try:
+                stream, filename = await self._afile_to_stream(instance)
+                yield "file", "file", stream, filename, "file"
+            except Exception as exc:
+                self._warn_skip("File", str(getattr(instance, "filename", instance)), exc)
+            return
+
+        model_fields = getattr(instance.__class__, "model_fields", None) or {}
+        if not model_fields:
+            return
+
+        for name, info in model_fields.items():
+            if fields_set is not None and name not in fields_set:
+                continue
+            val = getattr(instance, name, None)
+            try:
+                payload = await self._avalue_to_stream(val)
+                if payload is None:
+                    continue
+                stream, filename, kind = payload
+                title = getattr(info, "title", None) or name
+                yield name, title, stream, filename, kind
+            except Exception as exc:
+                self._warn_skip(instance.__class__.__name__, name, exc)
+
+    def _iter_values(
+        self,
+        instance: Any,
+        fields_set: set[str] | None,
+    ) -> Iterable[tuple[str, str | None, IO[Any], str | None, str]]:
+        if isinstance(instance, File):
+            try:
+                stream, filename = self._file_to_stream(instance)
+                return [("file", "file", stream, filename, "file")]
+            except Exception as exc:
+                self._warn_skip("File", str(getattr(instance, "filename", instance)), exc)
+                return []
+
+        model_fields = getattr(instance.__class__, "model_fields", None) or {}
+        if not model_fields:
+            return []
+
+        items: list[tuple[str, str | None, IO[Any], str | None, str]] = []
+        for name, info in model_fields.items():
+            if fields_set is not None and name not in fields_set:
+                continue
+            val = getattr(instance, name, None)
+            try:
+                payload = self._value_to_stream(val)
+                if payload is None:
+                    continue
+                stream, filename, kind = payload
+                title = getattr(info, "title", None) or name
+                items.append((name, title, stream, filename, kind))
+            except Exception as exc:
+                self._warn_skip(instance.__class__.__name__, name, exc)
+        return items
+
+    def _build_source(
+        self,
+        object_class: str,
+        object_id: str,
+        field_name: str,
+        field_title: str | None,
+        filename: str | None,
+    ) -> IngestionSource:
+        meta = dict(self.base_metadata)
+        meta.setdefault("field", field_name)
+        if field_title:
+            meta.setdefault("field_title", field_title)
+        if filename:
+            meta.setdefault("filename", filename)
+        return IngestionSource(
+            object_class=object_class,
+            object_id=object_id,
+            tags=list(self.base_tags),
+            metadata=meta,
+        )
+
+    def _build_tags(self, extra: Iterable[str] | None, field_name: str, kind: str) -> list[str]:
+        tags = list(extra or [])
+        tags.append(kind)
+        tags.append(f"field:{field_name}")
+        return tags
+
+    def _resolve_link(self, instance: Any) -> tuple[str, str]:
+        cls = instance.__class__.__name__
+        oid = getattr(instance, "object_id", None)
+        if oid is None:
+            oid = getattr(instance, "id", None)
+        if oid is None:
+            oid = id(instance)
+        return cls, str(oid)
+
+    def _warn_skip(self, object_class: str, field_name: str, exc: Exception) -> None:
+        self.logger.warning("Skipping %s.%s: %s", object_class, field_name, exc)
+
+    def _materialize(self, objects: Any) -> Iterable[Any]:
+        if isinstance(objects, QuerySetBase):
+            return objects.execute()  # type: ignore[attr-defined]
+        return list(objects)
+
+    async def _amaterialize(self, objects: Any) -> AsyncIterator[Any]:
+        if isinstance(objects, QuerySetBase):
+            result = await objects.aexecute()  # type: ignore[attr-defined]
+            for item in result:
+                yield item
+            return
+
+        for item in objects:
+            yield item
+
+    def _file_to_stream(self, file_obj: File) -> tuple[IO[Any], str | None]:
+        content = file_obj.read_bytes()
+        return io.BytesIO(content), getattr(file_obj, "filename", None)
+
+    async def _afile_to_stream(self, file_obj: File) -> tuple[IO[Any], str | None]:
+        content = await file_obj.aread_bytes()
+        return io.BytesIO(content), getattr(file_obj, "filename", None)
+
+    def _value_to_stream(self, val: Any) -> tuple[IO[Any], str | None, str] | None:
+        if val is None:
+            return None
+        if isinstance(val, File):
+            stream, filename = self._file_to_stream(val)
+            return stream, filename, "file"
+        if isinstance(val, (bytes, bytearray)):
+            return io.BytesIO(val), None, "file"
+        if isinstance(val, str):
+            return io.StringIO(val), None, "text"
+        return None
+
+    async def _avalue_to_stream(self, val: Any) -> tuple[IO[Any], str | None, str] | None:
+        if val is None:
+            return None
+        if isinstance(val, File):
+            stream, filename = await self._afile_to_stream(val)
+            return stream, filename, "file"
+        if isinstance(val, (bytes, bytearray)):
+            return io.BytesIO(val), None, "file"
+        if isinstance(val, str):
+            return io.StringIO(val), None, "text"
+        return None

amsdal_ml/ml_ingesting/pipeline.py

@@ -0,0 +1,131 @@
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import IO
+from typing import Any
+from typing import Generic
+from typing import TypeVar
+
+from amsdal_ml.ml_ingesting.embedders.embedder import Embedder
+from amsdal_ml.ml_ingesting.embedding_data import EmbeddingData
+from amsdal_ml.ml_ingesting.loaders.loader import Loader
+from amsdal_ml.ml_ingesting.pipeline_interface import IngestionPipeline
+from amsdal_ml.ml_ingesting.processors.cleaner import Cleaner
+from amsdal_ml.ml_ingesting.splitters.splitter import Splitter
+from amsdal_ml.ml_ingesting.stores.store import EmbeddingStore
+from amsdal_ml.ml_ingesting.types import IngestionSource
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+
+LoaderT = TypeVar('LoaderT', bound=Loader)
+
+
+class DefaultIngestionPipeline(IngestionPipeline, Generic[LoaderT]):
+    loader: LoaderT
+
+    def __init__(
+        self,
+        *,
+        loader: LoaderT,
+        cleaner: Cleaner,
+        splitter: Splitter,
+        embedder: Embedder,
+        store: EmbeddingStore,
+    ) -> None:
+        self.loader = loader
+        self.cleaner = cleaner
+        self.splitter = splitter
+        self.embedder = embedder
+        self.store = store
+
+    def _combine_tags(self, base: Iterable[str] | None, extra: Iterable[str] | None) -> list[str]:
+        out: list[str] = []
+        for tag in list(base or []) + list(extra or []):
+            if tag not in out:
+                out.append(tag)
+        return out
+
+    def _ensure_source(self, source: IngestionSource | None) -> IngestionSource:
+        if source is None:
+            msg = 'source is required for ingestion pipeline'
+            raise RuntimeError(msg)
+        return source
+
+    def _merge_metadata(
+        self,
+        source_meta: dict[str, Any],
+        doc_meta: dict[str, Any],
+        filename: str | None = None,
+    ) -> dict[str, Any]:
+        merged = {**source_meta, **doc_meta}
+        if filename and 'filename' not in merged:
+            merged['filename'] = filename
+        return merged
+
+    def _embed_chunks(self, chunks, tags: list[str], base_metadata: dict[str, Any]) -> list[EmbeddingData]:
+        embeddings: list[EmbeddingData] = []
+        for idx, chunk in enumerate(chunks):
+            vector = self.embedder.embed(chunk.text)
+            merged_tags = self._combine_tags(tags, chunk.tags)
+            metadata = {**base_metadata, **dict(chunk.metadata)}
+            embeddings.append(
+                EmbeddingData(
+                    chunk_index=idx,
+                    raw_text=chunk.text,
+                    embedding=vector,
+                    tags=merged_tags,
+                    metadata=metadata,
+                )
+            )
+        return embeddings
+
+    async def _aembed_chunks(self, chunks, tags: list[str], base_metadata: dict[str, Any]) -> list[EmbeddingData]:
+        embeddings: list[EmbeddingData] = []
+        for idx, chunk in enumerate(chunks):
+            vector = await self.embedder.aembed(chunk.text)
+            merged_tags = self._combine_tags(tags, chunk.tags)
+            metadata = {**base_metadata, **dict(chunk.metadata)}
+            embeddings.append(
+                EmbeddingData(
+                    chunk_index=idx,
+                    raw_text=chunk.text,
+                    embedding=vector,
+                    tags=merged_tags,
+                    metadata=metadata,
+                )
+            )
+        return embeddings
+
+    def run(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        tags: Iterable[str] | None = None,
+        source: IngestionSource | None = None,
+    ) -> list[Any]:
+        src = self._ensure_source(source)
+        doc = self.loader.load(file, filename=filename, metadata=src.metadata)
+        base_metadata = self._merge_metadata(src.metadata, doc.metadata, filename)
+        cleaned = self.cleaner.clean(LoadedDocument(pages=doc.pages, metadata=base_metadata))
+        chunks = self.splitter.split(cleaned)
+        merged_tags = self._combine_tags(src.tags, tags)
+        embeddings = self._embed_chunks(chunks, merged_tags, base_metadata)
+        return self.store.save(embeddings, source=src)
+
+    async def arun(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        tags: Iterable[str] | None = None,
+        source: IngestionSource | None = None,
+    ) -> list[Any]:
+        src = self._ensure_source(source)
+        doc = await self.loader.aload(file, filename=filename, metadata=src.metadata)
+        base_metadata = self._merge_metadata(src.metadata, doc.metadata, filename)
+        cleaned = await self.cleaner.aclean(LoadedDocument(pages=doc.pages, metadata=base_metadata))
+        chunks = await self.splitter.asplit(cleaned)
+        merged_tags = self._combine_tags(src.tags, tags)
+        embeddings = await self._aembed_chunks(chunks, merged_tags, base_metadata)
+        return await self.store.asave(embeddings, source=src)
+

amsdal_ml/ml_ingesting/pipeline_interface.py

@@ -0,0 +1,31 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+from collections.abc import Iterable
+from typing import IO
+from typing import Any
+
+from amsdal_ml.ml_ingesting.types import IngestionSource
+
+
+class IngestionPipeline(ABC):
+    @abstractmethod
+    def run(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        tags: Iterable[str] | None = None,
+        source: IngestionSource | None = None,
+    ) -> list[Any]: ...
+
+    @abstractmethod
+    async def arun(
+        self,
+        file: IO[Any],
+        *,
+        filename: str | None = None,
+        tags: Iterable[str] | None = None,
+        source: IngestionSource | None = None,
+    ) -> list[Any]: ...

amsdal_ml/ml_ingesting/processors/__init__.py

@@ -0,0 +1,4 @@
+from amsdal_ml.ml_ingesting.processors.cleaner import Cleaner
+from amsdal_ml.ml_ingesting.processors.text_cleaner import TextCleaner
+
+__all__ = ['Cleaner', 'TextCleaner']

amsdal_ml/ml_ingesting/processors/cleaner.py

@@ -0,0 +1,14 @@
+from __future__ import annotations
+
+from abc import ABC
+from abc import abstractmethod
+
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+
+
+class Cleaner(ABC):
+    @abstractmethod
+    def clean(self, doc: LoadedDocument) -> LoadedDocument: ...
+
+    @abstractmethod
+    async def aclean(self, doc: LoadedDocument) -> LoadedDocument: ...

amsdal_ml/ml_ingesting/processors/text_cleaner.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import asyncio
+import re
+
+from amsdal_ml.ml_ingesting.processors.cleaner import Cleaner
+from amsdal_ml.ml_ingesting.types import LoadedDocument
+from amsdal_ml.ml_ingesting.types import LoadedPage
+
+_CONTROL_CHARS = re.compile(r'[\x00-\x08\x0b\x0c\x0e-\x1f\x7f]')
+_MULTI_SPACE = re.compile(r'\s{2,}')
+
+
+def _normalize_text(text: str) -> str:
+    text = _CONTROL_CHARS.sub(' ', text)
+    text = text.replace('\r', ' ').replace('\t', ' ')
+    text = _MULTI_SPACE.sub(' ', text)
+    return text.strip()
+
+
+class TextCleaner(Cleaner):
+    def __init__(self, *, drop_empty_pages: bool = True) -> None:
+        self.drop_empty_pages = drop_empty_pages
+
+    def clean(self, doc: LoadedDocument) -> LoadedDocument:
+        pages: list[LoadedPage] = []
+        for page in doc.pages:
+            cleaned = _normalize_text(page.text)
+            if not cleaned and self.drop_empty_pages:
+                continue
+            pages.append(
+                LoadedPage(
+                    page_number=page.page_number,
+                    text=cleaned,
+                    metadata=dict(page.metadata),
+                )
+            )
+        meta = dict(doc.metadata)
+        return LoadedDocument(pages=pages, metadata=meta)
+
+    async def aclean(self, doc: LoadedDocument) -> LoadedDocument:
+        return await asyncio.to_thread(self.clean, doc)

amsdal_ml/ml_ingesting/splitters/__init__.py

@@ -0,0 +1,4 @@
+from amsdal_ml.ml_ingesting.splitters.splitter import Splitter
+from amsdal_ml.ml_ingesting.splitters.token_splitter import TokenSplitter
+
+__all__ = ['Splitter', 'TokenSplitter']