retrievalbase 1.0.0__py3-none-any.whl

retrievalbase/connector/__init__.py

@@ -0,0 +1,69 @@
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+
+from retrievalbase.connector.settings import DatasetConnectorSettings
+from retrievalbase.mixins import FromConfigMixin
+from retrievalbase.types import TCDatasetConnector as TCDatasetConnector
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import Dataset, TextDataset
+
+_logger = logging.getLogger(__name__)
+
+
+class DatasetConnector[TCDatasetConnector: DatasetConnectorSettings](
+    FromConfigMixin[TCDatasetConnector],
+    ABC,
+):
+    def __init__(self, config: TCDatasetConnector):
+        super().__init__(config)
+
+        _logger.info(
+            f"Initializing dataset connector | class={self.__class__.__name__} | module={self.__class__.__module__}"
+        )
+
+    @abstractmethod
+    def _load(self) -> pl.DataFrame | pl.LazyFrame:
+        """
+        Load raw data as Polars DataFrame or LazyFrame.
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def to(self, ds: "Dataset[Any]") -> None:
+        raise NotImplementedError()
+
+    def load(self) -> "Dataset[pl.DataFrame | pl.LazyFrame]":
+        from retrievalbase.dataset.polars import PolarsDataset
+
+        _logger.info(f"Loading dataset | connector={self.__class__.__name__}")
+
+        df = self._load()
+        self._log_polars_info(df)
+
+        return PolarsDataset.from_polars(df)
+
+    def load_text(self) -> "TextDataset[pl.DataFrame | pl.LazyFrame]":
+        from retrievalbase.dataset.polars import PolarsTextDataset
+
+        _logger.info(f"Loading text dataset | connector={self.__class__.__name__}")
+
+        df = self._load()
+        self._log_polars_info(df)
+
+        return PolarsTextDataset.from_polars(df)
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+
+    def _log_polars_info(self, df: pl.DataFrame | pl.LazyFrame) -> None:
+        """
+        Log dataset structure without forcing materialization.
+        """
+        if isinstance(df, pl.DataFrame):
+            schema = df.schema
+            _logger.info(f"Loaded DataFrame | columns={len(schema)} | schema={list(schema.items())}")

retrievalbase/connector/minio.py

@@ -0,0 +1,45 @@
+import io
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+from minio import Minio
+
+from retrievalbase.connector import DatasetConnector
+from retrievalbase.connector.settings import MinioDatasetConnectorSettings
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import Dataset
+
+
+class MinioDatasetConnector(DatasetConnector[MinioDatasetConnectorSettings]):
+    def __init__(self, config: MinioDatasetConnectorSettings):
+        super().__init__(config)
+        self.client = Minio(
+            self.config.endpoint.replace("http://", "").replace("https://", ""),
+            access_key=self.config.access_key.get_secret_value(),
+            secret_key=self.config.secret_key.get_secret_value(),
+            secure=self.config.endpoint.startswith("https://"),
+        )
+
+    def _load(self) -> pl.DataFrame | pl.LazyFrame:
+        response = self.client.get_object(self.config.bucket, self.config.key)
+        try:
+            buffer = io.BytesIO(response.read())
+        finally:
+            response.close()
+            response.release_conn()
+        df = pl.read_parquet(buffer)
+        return df
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        df = ds.polars
+        buffer = io.BytesIO()
+        df.write_parquet(buffer)
+        buffer.seek(0)
+        self.client.put_object(
+            bucket_name=self.config.bucket,
+            object_name=self.config.key,
+            data=buffer,
+            length=buffer.getbuffer().nbytes,
+            content_type="application/octet-stream",
+        )

retrievalbase/connector/parquet.py

@@ -0,0 +1,20 @@
+from typing import TYPE_CHECKING, Any
+
+import polars as pl
+
+from retrievalbase.connector import DatasetConnector
+from retrievalbase.connector.settings import ParquetDatasetConnectorSettings
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import Dataset
+
+
+class ParquetDatasetConnector(DatasetConnector[ParquetDatasetConnectorSettings]):
+    def __init__(self, config: ParquetDatasetConnectorSettings):
+        super().__init__(config)
+
+    def _load(self) -> pl.DataFrame | pl.LazyFrame:
+        return pl.scan_parquet(self.config.path) if self.config.lazy else pl.read_parquet(self.config.path)
+
+    def to(self, ds: "Dataset[Any]") -> None:
+        ds.polars.write_parquet(self.config.path)

retrievalbase/connector/settings.py

@@ -0,0 +1,22 @@
+from pydantic import SecretStr
+from pydantic_settings import SettingsConfigDict
+
+from retrievalbase.settings import FromConfigMixinSettings
+
+
+class DatasetConnectorSettings(FromConfigMixinSettings):
+    pass
+
+
+class ParquetDatasetConnectorSettings(DatasetConnectorSettings):
+    path: str
+    lazy: bool
+
+
+class MinioDatasetConnectorSettings(DatasetConnectorSettings):
+    endpoint: str
+    bucket: str
+    key: str
+    access_key: SecretStr
+    secret_key: SecretStr
+    model_config = SettingsConfigDict(env_prefix="MINIO_", extra="ignore")

retrievalbase/constants.py

@@ -0,0 +1 @@
+CONFIG_PATH = "/config/config.yaml"

retrievalbase/dataset/__init__.py

@@ -0,0 +1,146 @@
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Iterator
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Literal,
+    Self,
+    cast,
+    overload,
+)
+
+import polars as pl
+from polars._typing import ColumnNameOrSelector, IntoExpr, IntoExprColumn
+
+from retrievalbase.dataset.mixins import TextDatasetMixin
+from retrievalbase.types import TDataset as TDataset
+from retrievalbase.utils import _get_minio_connector, _get_parquet_connector
+
+if TYPE_CHECKING:
+    pass
+
+
+class Dataset[TDataset](ABC):
+    def __init__(self, service: TDataset):
+        super().__init__()
+        self.service = service
+        self._polars: pl.DataFrame | None = None
+        self._lazy_polars: pl.LazyFrame | None = None
+
+    @classmethod
+    @abstractmethod
+    def from_polars(cls, df: pl.DataFrame | pl.LazyFrame) -> Self:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_polars(self) -> pl.DataFrame:
+        raise NotImplementedError
+
+    @abstractmethod
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        raise NotImplementedError
+
+    @classmethod
+    def from_parquet(
+        cls,
+        path: str,
+        *,
+        lazy: bool = True,
+    ) -> "Dataset[pl.DataFrame | pl.LazyFrame]":
+        return _get_parquet_connector(path, lazy=lazy).load()
+
+    def to_parquet(self, path: str, *, lazy: bool = True) -> None:
+        return _get_parquet_connector(path, lazy=lazy).to(self)
+
+    @classmethod
+    def from_minio(
+        cls,
+        bucket: str,
+        key: str,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+    ) -> "Dataset[pl.DataFrame | pl.LazyFrame]":
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).load()
+
+    def to_minio(self, bucket: str, key: str, endpoint: str, access_key: str, secret_key: str) -> None:
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).to(self)
+
+    @property
+    def polars(self) -> pl.DataFrame:
+        if self._polars is None:
+            self._polars = self.to_polars()
+        return self._polars
+
+    @property
+    def lazy_polars(self) -> pl.LazyFrame:
+        if self._lazy_polars is None:
+            self._lazy_polars = self.to_lazy_polars()
+        return self._lazy_polars
+
+    @property
+    def polars_shape(self) -> tuple[int, int]:
+        return self.polars.shape
+
+    def with_columns(
+        self,
+        *exprs: IntoExpr | Iterable[IntoExpr],
+        **named_exprs: IntoExpr,
+    ) -> Self:
+        df = self.polars.with_columns(*exprs, **named_exprs)
+        return self.__class__.from_polars(df)
+
+    def filter(
+        self,
+        *predicates: (IntoExprColumn | Iterable[IntoExprColumn] | bool | list[bool]),
+        **constraints: Any,
+    ) -> Self:
+        return self.__class__.from_polars(self.polars.filter(*predicates, **constraints))
+
+    def drop(
+        self,
+        *columns: ColumnNameOrSelector | Iterable[ColumnNameOrSelector],
+        strict: bool = True,
+    ) -> Self:
+        return self.__class__.from_polars(self.polars.drop(*columns, strict=strict))
+
+    def __len__(self) -> int:
+        return self.polars_shape[0]
+
+    @overload
+    def iter_rows(self, *, named: Literal[False] = ..., buffer_size: int = ...) -> Iterator[tuple[Any, ...]]: ...
+
+    @overload
+    def iter_rows(self, *, named: Literal[True], buffer_size: int = ...) -> Iterator[dict[str, Any]]: ...
+
+    def iter_rows(
+        self, *, named: Literal[False, True] = False, buffer_size: int = 512
+    ) -> Iterator[tuple[Any, ...]] | Iterator[dict[str, Any]]:
+        df_stream: pl.DataFrame = cast(pl.DataFrame, self.lazy_polars.collect())
+        return df_stream.iter_rows(named=named, buffer_size=buffer_size)  # type: ignore[no-any-return]
+
+
+class TextDataset[TDataset](Dataset[TDataset], TextDatasetMixin[TDataset]):
+    def __init__(self, service: TDataset):
+        super().__init__(service)
+        self._validate_schema()
+
+    @classmethod
+    def from_parquet(
+        cls,
+        path: str,
+        *,
+        lazy: bool = True,
+    ) -> "TextDataset[pl.DataFrame | pl.LazyFrame]":
+        return _get_parquet_connector(path, lazy=lazy).load_text()
+
+    @classmethod
+    def from_minio(
+        cls,
+        bucket: str,
+        key: str,
+        endpoint: str,
+        access_key: str,
+        secret_key: str,
+    ) -> "TextDataset[pl.DataFrame | pl.LazyFrame]":
+        return _get_minio_connector(bucket, key, endpoint, access_key, secret_key).load_text()

retrievalbase/dataset/hf.py

@@ -0,0 +1,49 @@
+from abc import ABC
+from typing import Any
+
+import polars as pl
+from datasets import Dataset as HFDataset
+
+from retrievalbase.connector import DatasetConnector
+from retrievalbase.dataset import TextDataset
+from retrievalbase.dataset.settings import HuggingFaceDatasetAdaptaterSettings
+from retrievalbase.exceptions import DatasetSchemaError
+from retrievalbase.mixins import FromConfigMixin
+from retrievalbase.utils import build_schema, extract_schema_columns, load_class
+
+
+class HuggingFaceDatasetAdaptater[TCHFDatasetAdaptater: HuggingFaceDatasetAdaptaterSettings[Any]](
+    FromConfigMixin[TCHFDatasetAdaptater],
+    ABC,
+):
+    def __init__(self, config: TCHFDatasetAdaptater) -> None:
+        super().__init__(config)
+        self._dataset: TextDataset[pl.DataFrame | pl.LazyFrame] | None = None
+
+    @property
+    def dataset(self) -> TextDataset:
+        if self._dataset is None:
+            connector: DatasetConnector[Any] = load_class(self.config.connector.module_path).from_config(
+                self.config.connector
+            )
+
+            self._dataset = connector.load_text()
+
+        return self._dataset
+
+    def to_hf(self) -> HFDataset:
+        df = self.dataset.polars
+        # schema mapping
+        schema = self.config.columns_mapping
+        # validation (top-level)
+        available_columns = df.columns
+        flattened_columns = extract_schema_columns(df.schema)
+        missing_roots = {path.split(".")[0] for path in schema.values() if path.split(".")[0] not in available_columns}
+        if missing_roots:
+            raise DatasetSchemaError(
+                missing_columns=missing_roots,
+                available_columns=available_columns,
+                flattened_columns=flattened_columns,
+            )
+        df = build_schema(df, schema)
+        return HFDataset.from_dict(df.to_dict(as_series=False))

retrievalbase/dataset/mixins.py

@@ -0,0 +1,108 @@
+import json
+from collections.abc import Iterator
+from pathlib import Path
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    Protocol,
+    runtime_checkable,
+)
+
+import polars as pl
+from langchain_core.documents import Document
+
+from retrievalbase.types import TDataset as TDataset
+
+if TYPE_CHECKING:
+    from retrievalbase.dataset import TextDataset
+
+
+@runtime_checkable
+class SupportsPolarsDataset(Protocol[TDataset]):
+    @property
+    def polars(self) -> pl.DataFrame: ...
+
+    def to_polars(self) -> pl.DataFrame: ...
+
+    @classmethod
+    def from_polars(cls, df: pl.DataFrame | pl.LazyFrame) -> "TextDataset[TDataset]":
+        raise NotImplementedError
+
+
+class TextDatasetMixin[TDataset](SupportsPolarsDataset[TDataset]):
+    REQUIRED_COLUMNS = {"page_content", "metadata"}
+
+    def _validate_schema(self) -> None:
+        df = self.polars
+        missing = self.REQUIRED_COLUMNS - set(df.columns)
+        if missing:
+            raise ValueError(
+                f"{self.__class__.__name__} requires columns {self.REQUIRED_COLUMNS}, but missing {missing}"
+            )
+
+    def iter_documents(self) -> Iterator[Document]:
+        df = self.to_polars()
+        for row in df.iter_rows(named=True):
+            yield Document(
+                page_content=row["page_content"],
+                metadata=row["metadata"],
+            )
+
+    @classmethod
+    def from_records(
+        cls,
+        records: list[tuple[str, dict[str, Any]]],
+    ) -> "TextDataset[TDataset]":
+        if not records:
+            raise ValueError("records must be non-empty")
+        df = pl.DataFrame(
+            {
+                "page_content": [text for text, _ in records],
+                "metadata": [meta for _, meta in records],
+            }
+        )
+        return cls.from_polars(df)
+
+    def dump_documents(
+        self,
+        out_dir: str,
+        prefix: str = "doc",
+        ext: str = ".md",
+        encoding: str = "utf-8",
+    ) -> None:
+        out_dir_path = Path(out_dir)
+        out_dir_path.mkdir(parents=True, exist_ok=True)
+
+        for i, row in enumerate(self.polars.iter_rows(named=True)):
+            path = out_dir_path / f"{prefix}_{i:05d}{ext}"
+            page_content = row.get("page_content", "")
+            metadata = row.get("metadata", {})
+            try:
+                metadata_str = json.dumps(metadata, indent=2, ensure_ascii=False)
+            except TypeError:
+                metadata_str = str(metadata)
+            content = f"###PAGE_CONTENT###\n{page_content}\n\n###METADATA###\n{metadata_str}\n"
+            path.write_text(content, encoding=encoding)
+
+    def to_langchain_documents(
+        self,
+    ) -> list[Document]:
+        docs: list[Document] = []
+        for row in self.polars.iter_rows(named=True):
+            docs.append(
+                Document(
+                    page_content=row["page_content"],
+                    metadata=row["metadata"],
+                )
+            )
+        return docs
+
+    @classmethod
+    def from_documents(cls, docs: list[Document]) -> "TextDataset[TDataset]":
+        df = pl.DataFrame(
+            {
+                "page_content": [d.page_content for d in docs],
+                "metadata": [d.metadata or {} for d in docs],
+            }
+        )
+        return cls.from_polars(df)

retrievalbase/dataset/polars.py

@@ -0,0 +1,43 @@
+from typing import Self
+
+import polars as pl
+
+from retrievalbase.dataset import Dataset, TextDataset
+
+
+class PolarsDataset(Dataset[pl.DataFrame | pl.LazyFrame]):
+    def __init__(self, service: pl.DataFrame | pl.LazyFrame):
+        super().__init__(service)
+
+    @classmethod
+    def from_polars(cls, df: pl.DataFrame | pl.LazyFrame) -> Self:
+        return cls(df)
+
+    def to_polars(self) -> pl.DataFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service.collect(background=False)  # ty: ignore[invalid-return-type]
+        return self.service
+
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service
+        return self.service.lazy()
+
+
+class PolarsTextDataset(TextDataset[pl.DataFrame | pl.LazyFrame]):
+    def __init__(self, service: pl.DataFrame | pl.LazyFrame):
+        super().__init__(service)
+
+    @classmethod
+    def from_polars(cls, df: pl.DataFrame | pl.LazyFrame) -> Self:
+        return cls(df)
+
+    def to_polars(self) -> pl.DataFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service.collect(background=False)  # ty: ignore[invalid-return-type]
+        return self.service
+
+    def to_lazy_polars(self) -> pl.LazyFrame:
+        if isinstance(self.service, pl.LazyFrame):
+            return self.service
+        return self.service.lazy()

retrievalbase/dataset/preprocess/__init__.py

@@ -0,0 +1,29 @@
+from abc import ABC, abstractmethod
+from typing import Any
+
+from retrievalbase.dataset.polars import TextDataset
+from retrievalbase.dataset.settings import TextPreprocessorSettings, TokenCounterSettings
+from retrievalbase.mixins import FromConfigMixin
+from retrievalbase.types import TCTextPreprocessor as TCTextPreprocessor
+from retrievalbase.types import TCTokenCounter as TCTokenCounter
+
+
+class TextPreprocessor[TCTextPreprocessor: TextPreprocessorSettings](FromConfigMixin[TCTextPreprocessor], ABC):
+    def __init__(self, config: TCTextPreprocessor) -> None:
+        self.config = config
+
+    @abstractmethod
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        raise NotImplementedError
+
+
+class TokenCounter[TCTokenCounter: TokenCounterSettings](FromConfigMixin[TCTokenCounter], ABC):
+    def __init__(self, config: TCTokenCounter) -> None:
+        self.config = config
+
+    @abstractmethod
+    def count(self, text: str) -> int:
+        raise NotImplementedError()
+
+    def count_batch(self, texts: list[str]) -> list[int]:
+        return [self.count(t) for t in texts]

retrievalbase/dataset/preprocess/preprocess.py

@@ -0,0 +1,96 @@
+from typing import Any
+
+import polars as pl
+
+from retrievalbase.dataset import TextDataset
+from retrievalbase.dataset.preprocess import TextPreprocessor, TokenCounter
+from retrievalbase.dataset.settings import (
+    MaxTokenFilterSettings,
+    MinTokenFilterSettings,
+    PreprocessPipelineSettings,
+    QuantileTokenFilterSettings,
+    SigmaBandTokenFilterSettings,
+)
+from retrievalbase.utils import load_class
+
+
+class MinTokenFilter(TextPreprocessor[MinTokenFilterSettings[Any]]):
+    def __init__(self, config: MinTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        return ds.__class__.from_polars(
+            ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+            .filter(pl.col("n_tokens") >= self.config.min_tokens)
+            .drop("n_tokens")
+        )
+
+
+class MaxTokenFilter(TextPreprocessor[MaxTokenFilterSettings[Any]]):
+    def __init__(self, config: MaxTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        return ds.__class__.from_polars(
+            ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+            .filter(pl.col("n_tokens") <= self.config.max_tokens)
+            .drop("n_tokens")
+        )
+
+
+class QuantileTokenFilter(TextPreprocessor[QuantileTokenFilterSettings[Any]]):
+    def __init__(self, config: QuantileTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        df = ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+        cutoff = df.select(pl.col("n_tokens").quantile(self.config.q)).item()
+        return ds.__class__.from_polars(df.filter(pl.col("n_tokens") <= cutoff).drop("n_tokens"))
+
+
+class SigmaBandTokenFilter(TextPreprocessor[SigmaBandTokenFilterSettings[Any]]):
+    def __init__(self, config: SigmaBandTokenFilterSettings[Any]):
+        super().__init__(config)
+        self.token_counter: TokenCounter[Any] = load_class(self.config.token_counter.module_path).from_config(
+            self.config.token_counter
+        )
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        df = ds.polars.with_columns(pl.col("page_content").map_elements(self.token_counter.count).alias("n_tokens"))
+        mu, sigma = df.select(
+            pl.col("n_tokens").mean().alias("mu"),
+            pl.col("n_tokens").std().alias("sigma"),
+        ).row(0)
+        return ds.__class__.from_polars(
+            df.filter(pl.col("n_tokens").is_between(mu - self.config.z * sigma, mu + self.config.z * sigma)).drop(
+                "n_tokens"
+            )
+        )
+
+
+class PreprocessPipeline(TextPreprocessor[PreprocessPipelineSettings]):
+    def __init__(self, config: PreprocessPipelineSettings):
+        super().__init__(config)
+        self.steps = self._load_steps()
+
+    def _load_steps(self) -> list[TextPreprocessor[Any]]:
+        steps = []
+        for step_config in self.config.steps:
+            step: TextPreprocessor[Any] = load_class(step_config.module_path)(step_config)
+            steps.append(step)
+        return steps
+
+    def apply(self, ds: TextDataset[Any]) -> TextDataset[Any]:
+        current = ds
+        for step in self.steps:
+            current = step.apply(current)
+        return current

retrievalbase/dataset/preprocess/token_counter.py

@@ -0,0 +1,41 @@
+from retrievalbase.dataset.preprocess import TokenCounter
+from retrievalbase.dataset.settings import (
+    HeuristicTokenCounterSettings,
+    HuggingFaceTokenCounterSettings,
+)
+
+
+class HeuristicTokenCounter(TokenCounter[HeuristicTokenCounterSettings]):
+    def __init__(self, config: HeuristicTokenCounterSettings) -> None:
+        super().__init__(config)
+
+    def count(self, text: str) -> int:
+        return max(1, int(len(text) / self.config.chars_per_token))
+
+
+class HuggingFaceTokenCounter(TokenCounter[HuggingFaceTokenCounterSettings]):
+    def __init__(self, config: HuggingFaceTokenCounterSettings) -> None:
+        super().__init__(config)
+        try:
+            from transformers import AutoTokenizer
+        except ImportError as e:
+            raise ImportError(
+                "HuggingFaceTokenCounter requires transformers.\nInstall with: pip install retrievalbase[transformers]"
+            ) from e
+
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            self.config.name,
+            revision=self.config.revision,
+            trust_remote_code=True,
+        )
+
+    def count(self, text: str) -> int:
+        return len(
+            self.tokenizer(  # ty: ignore[call-non-callable]
+                text,
+                add_special_tokens=self.config.add_special_tokens,
+                truncation=False,
+                return_attention_mask=False,
+                return_token_type_ids=False,
+            )["input_ids"]
+        )