diversify-text 0.1.1__py3-none-any.whl

diversify_text/__init__.py

@@ -0,0 +1,24 @@
+"""diversify-text -- generate stylistic paraphrases of texts."""
+
+import logging
+
+from diversify_text.core import (
+    Diversifier,
+    diversify,
+)
+
+__all__ = [
+    "Diversifier",
+    "diversify",
+]
+
+# Configure a clean handler for the diversify logger so INFO/WARNING messages
+# are visible without requiring the user to set up logging themselves.
+_logger = logging.getLogger("diversify_text")
+_logger.setLevel(logging.INFO)
+_handler = logging.StreamHandler()
+_handler.setFormatter(logging.Formatter("%(levelname)s: %(message)s"))
+_logger.addHandler(_handler)
+# Prevent messages from bubbling up to the root logger (avoids duplicate output
+# if the user has already configured logging globally).
+_logger.propagate = False

diversify_text/_input.py

@@ -0,0 +1,267 @@
+"""Input resolution for diversify.
+
+Converts the many input forms users can provide (single string, list,
+generator, CSV/TSV/TXT file path) into a uniform ``Iterator[str]`` plus
+an :class:`InputContext` that describes the source.
+"""
+
+from __future__ import annotations
+
+import csv
+import logging
+from collections.abc import Iterable, Iterator
+from dataclasses import dataclass
+from enum import Enum, auto
+from pathlib import Path
+from typing import Union
+
+
+# ------------------------------------------------------------------
+# Type alias
+# ------------------------------------------------------------------
+
+TextInput = Union[str, "list[str]", "Iterable[str]"]
+
+_log = logging.getLogger(__name__)
+
+
+# ------------------------------------------------------------------
+# Input kind & context
+# ------------------------------------------------------------------
+
+
+class InputKind(Enum):
+    """Discriminator for how the user provided input."""
+
+    SINGLE_STR = auto()
+    LIST = auto()
+    ITERABLE = auto()
+    FILE_CSV = auto()
+    FILE_TSV = auto()
+    FILE_TXT = auto()
+
+
+@dataclass(frozen=True)
+class InputContext:
+    """Read-only metadata about the resolved input source."""
+
+    kind: InputKind
+    input_path: Path | None = None
+    text_column: str | None = None
+    total: int | None = None
+
+
+# ------------------------------------------------------------------
+# Public API
+# ------------------------------------------------------------------
+
+
+def resolve_input(
+    texts: TextInput,
+    text_column: str = "text",
+) -> tuple[Iterator[str], InputContext]:
+    """Convert any supported input into a lazy ``Iterator[str]`` plus metadata.
+
+    Parameters
+    ----------
+    texts : str | list[str] | Iterable[str]
+        A single text, a list of texts, a generator / iterable of texts,
+        or a path to a ``.csv``, ``.tsv``, or ``.txt`` file.
+    text_column : str
+        Column name to extract when *texts* points to a CSV/TSV file.
+
+    Returns
+    -------
+    (Iterator[str], InputContext)
+    """
+    # --- str: could be a file path or a single text ---
+    if isinstance(texts, str):
+        path = Path(texts)
+        suffix = path.suffix.lower()
+
+        if suffix == ".csv" and path.is_file():
+            _validate_csv_header(path, text_column, delimiter=",")
+            total = _count_file_lines(path) - 1  # subtract header row
+            return _iter_csv(path, text_column, delimiter=","), InputContext(
+                kind=InputKind.FILE_CSV,
+                input_path=path,
+                text_column=text_column,
+                total=total,
+            )
+
+        if suffix == ".tsv" and path.is_file():
+            _validate_csv_header(path, text_column, delimiter="\t")
+            total = _count_file_lines(path) - 1  # subtract header row
+            return _iter_csv(path, text_column, delimiter="\t"), InputContext(
+                kind=InputKind.FILE_TSV,
+                input_path=path,
+                text_column=text_column,
+                total=total,
+            )
+
+        if suffix == ".txt" and path.is_file():
+            total = _count_nonempty_lines(path)
+            _log.warning(
+                "TXT file input: each line is treated as a separate text "
+                "to diversify. Newlines are never part of a parsed text. "
+                "(%d non-empty lines found in %s)",
+                total,
+                path,
+            )
+            return _iter_txt_lines(path), InputContext(
+                kind=InputKind.FILE_TXT,
+                input_path=path,
+                total=total,
+            )
+
+        # Not a recognized file — treat as a single text.
+        return iter([texts]), InputContext(kind=InputKind.SINGLE_STR, total=1)
+
+    # --- list[str] ---
+    if isinstance(texts, list):
+        return iter(texts), InputContext(kind=InputKind.LIST, total=len(texts))
+
+    # --- generic Iterable[str] (generators, file handles, …) ---
+    if isinstance(texts, Iterable):
+        return iter(texts), InputContext(kind=InputKind.ITERABLE, total=None)
+
+    raise TypeError(
+        f"Unsupported input type {type(texts).__name__}. "
+        "Expected str, list[str], or Iterable[str]."
+    )
+
+
+# ------------------------------------------------------------------
+# Validation & counting (private, cheap first-pass helpers)
+# ------------------------------------------------------------------
+
+
+def _validate_csv_header(path: Path, text_column: str, delimiter: str) -> None:
+    """Check that *text_column* exists in the CSV/TSV header.
+
+    Opens the file, reads only the header row, then closes it.
+
+    Parameters
+    ----------
+    path : Path
+        Path to the CSV or TSV file.
+    text_column : str
+        Expected column name.
+    delimiter : str
+        Field separator — ``","`` for CSV, ``"\\t"`` for TSV.
+
+    Raises
+    ------
+    ValueError
+        If *text_column* is not found among the file's header fields.
+    """
+    with open(path, newline="", encoding="utf-8") as f:
+        reader = csv.DictReader(f, delimiter=delimiter)
+        if reader.fieldnames is None or text_column not in reader.fieldnames:
+            available = ", ".join(reader.fieldnames or [])
+            raise ValueError(
+                f"Column '{text_column}' not found in {path}. "
+                f"Available: {available}"
+            )
+
+
+def _count_file_lines(path: Path) -> int:
+    """Count the total number of lines in a file.
+
+    This is a cheap pass that only reads raw lines — no CSV parsing.
+    For CSV/TSV files, subtract 1 for the header to get the data-row
+    count.  The count may slightly overestimate if the file contains
+    multi-line quoted CSV fields, but that is acceptable for a progress
+    bar.
+
+    Parameters
+    ----------
+    path : Path
+        Path to the file.
+
+    Returns
+    -------
+    int
+        Total number of lines.
+    """
+    with open(path, encoding="utf-8") as f:
+        return sum(1 for _ in f)
+
+
+def _count_nonempty_lines(path: Path) -> int:
+    """Count non-empty lines in a file (skips blank / whitespace-only).
+
+    Parameters
+    ----------
+    path : Path
+        Path to the file.
+
+    Returns
+    -------
+    int
+        Number of non-empty lines.
+    """
+    with open(path, encoding="utf-8") as f:
+        return sum(1 for line in f if line.strip())
+
+
+# ------------------------------------------------------------------
+# Lazy file iterators (private)
+# ------------------------------------------------------------------
+
+
+def _iter_csv(path: Path, text_column: str, delimiter: str) -> Iterator[str]:
+    """Lazily yield *text_column* values from a CSV/TSV file.
+
+    Only one row is held in memory at a time.  The file handle is
+    closed automatically when the generator is exhausted or garbage
+    collected.
+
+    Parameters
+    ----------
+    path : Path
+        Path to the CSV or TSV file.
+    text_column : str
+        Name of the column that contains the texts to diversify.
+    delimiter : str
+        Field separator — ``","`` for CSV, ``"\\t"`` for TSV.
+
+    Yields
+    ------
+    str
+        The text value from each row.
+    """
+    f = open(path, newline="", encoding="utf-8")
+    try:
+        reader = csv.DictReader(f, delimiter=delimiter)
+        for row in reader:
+            yield row.get(text_column) or ""
+    finally:
+        f.close()
+
+
+def _iter_txt_lines(path: Path) -> Iterator[str]:
+    """Lazily yield non-empty, stripped lines from a ``.txt`` file.
+
+    Blank lines and whitespace-only lines are skipped.  Only one line
+    is held in memory at a time.  The file handle is closed
+    automatically when the generator is exhausted or garbage collected.
+
+    Parameters
+    ----------
+    path : Path
+        Path to the ``.txt`` file.
+
+    Yields
+    ------
+    str
+        Each non-empty line, with leading/trailing whitespace stripped.
+    """
+    f = open(path, encoding="utf-8")
+    try:
+        for line in f:
+            stripped = line.strip()
+            if stripped:
+                yield stripped
+    finally:
+        f.close()

diversify_text/_output.py

@@ -0,0 +1,234 @@
+"""Output path resolution and incremental writing for diversify.
+
+Decides *where* results go (in-memory vs. disk) and writes them in the
+appropriate format (Python list or JSONL).
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+from typing import IO, Any, Union
+
+from diversify_text._input import InputContext, InputKind
+
+_log = logging.getLogger(__name__)
+
+
+# ------------------------------------------------------------------
+# Type alias
+# ------------------------------------------------------------------
+
+DiversifyOutput = Union[list[dict], Path]
+
+
+# ------------------------------------------------------------------
+# Output path resolution
+# ------------------------------------------------------------------
+
+
+def resolve_output_path(
+    input_context: InputContext,
+    output_dir: str | Path | None = None,
+    output_name: str | None = None,
+) -> Path | None:
+    """Determine where output should be written, or ``None`` for in-memory.
+
+    The user controls *where* (directory) and *what name* (stem) to use,
+    but the **extension is always ``.jsonl``**.
+
+    Directory defaults (when *output_dir* is ``None``):
+
+    * ``SINGLE_STR`` / ``LIST`` → ``None`` (keep in memory, return as
+      ``list[dict]``).  These are small, known-size inputs from Python
+      code, so the caller typically wants results as Python objects.
+    * ``ITERABLE`` → current working directory.
+    * ``FILE_CSV`` / ``FILE_TSV`` / ``FILE_TXT`` → same directory as
+      the input file.
+
+    If *output_dir* is provided for ``SINGLE_STR`` / ``LIST``, results
+    are written to disk instead of being returned in memory.
+
+    Name defaults (when *output_name* is ``None``):
+
+    * ``FILE_CSV`` / ``FILE_TSV`` → ``<input_stem>_diversified``
+    * ``FILE_TXT`` → ``<input_stem>``
+    * Everything else → ``diversified_output``
+
+    Parameters
+    ----------
+    input_context : InputContext
+        Metadata produced by :func:`resolve_input`.
+    output_dir : str, Path, or None
+        Directory to write output files into.
+    output_name : str or None
+        Base filename (without extension).  The correct extension is
+        appended automatically.  If the name already contains an
+        extension it is **not** stripped — the correct extension is
+        appended after it — unless it already ends with ``.jsonl``.
+
+    Returns
+    -------
+    Path or None
+        ``None`` means in-memory mode; otherwise the path to write to.
+    """
+    # --- determine directory ---
+    if output_dir is not None:
+        directory = Path(output_dir)
+    elif input_context.kind in (InputKind.SINGLE_STR, InputKind.LIST):
+        # No output_dir and in-memory input → stay in-memory.
+        return None
+    elif input_context.kind == InputKind.ITERABLE:
+        # Iterable with no output_dir → default to current working directory.
+        directory = Path.cwd()
+    else:
+        # FILE_CSV, FILE_TSV, FILE_TXT → same directory as input file.
+        assert input_context.input_path is not None
+        directory = input_context.input_path.parent
+
+    # --- determine base name ---
+    if output_name is not None:
+        name = output_name
+    elif input_context.kind in (InputKind.FILE_CSV, InputKind.FILE_TSV):
+        assert input_context.input_path is not None
+        name = f"{input_context.input_path.stem}_diversified"
+    elif input_context.kind == InputKind.FILE_TXT:
+        assert input_context.input_path is not None
+        name = input_context.input_path.stem
+    else:
+        # ITERABLE, or LIST/SINGLE_STR with output_dir.
+        name = "diversified_output"
+
+    # --- build final path with the correct extension ---
+    if not name.endswith(".jsonl"):
+        name = f"{name}.jsonl"
+    result = directory / name
+
+    _log.info("Output will be written to %s", result)
+    return result
+
+
+# ------------------------------------------------------------------
+# Output writer
+# ------------------------------------------------------------------
+
+
+class OutputWriter:
+    """Incrementally writes diversify results to the right format.
+
+    Modes
+    -----
+    * **In-memory** (``output_path is None``): accumulates
+      ``list[dict]`` with keys ``"original"`` and ``"paraphrases"``.
+    * **JSONL** (``output_path is not None``): writes one JSON object
+      per line to a ``.jsonl`` file.
+    """
+
+    def __init__(
+        self,
+        input_context: InputContext,
+        n_styles: int,
+        output_path: Path | None,
+    ) -> None:
+        """Initialize the writer.
+
+        Parameters
+        ----------
+        input_context : InputContext
+            Metadata about the input source (kind, path, etc.).
+        n_styles : int
+            Number of paraphrase styles requested per text.
+        output_path : Path or None
+            Where to write results on disk.  ``None`` means results
+            are kept in memory and returned as ``list[dict]``.
+        """
+        self._input_context = input_context
+        self._n_styles = n_styles
+        self._output_path = output_path
+        # Open file handle — set by open() when writing to disk.
+        self._handle: IO[str] | None = None
+        # In-memory accumulator — used only when output_path is None.
+        self._accumulated: list[dict[str, Any]] = []
+
+    # --- lifecycle: open / write / close ---
+
+    def open(self) -> None:
+        """Open the file handle when writing to disk.
+
+        Must be called before :meth:`write_batch`.
+
+        * ``output_path is None`` — does nothing (in-memory mode).
+        * Otherwise — opens a single JSONL file for writing.
+        """
+        if self._output_path is None:
+            # In-memory mode: nothing to open.
+            return
+
+        # Disk mode: open a single JSONL file.
+        self._output_path.parent.mkdir(parents=True, exist_ok=True)
+        self._handle = open(self._output_path, "w", encoding="utf-8")
+
+    def write_batch(
+        self,
+        originals: list[str],
+        paraphrases_by_text: list[list[str]],
+    ) -> None:
+        """Append one batch of results.
+
+        Parameters
+        ----------
+        originals : list[str]
+            The original texts in this batch.
+        paraphrases_by_text : list[list[str]]
+            One inner list per original text, each containing *n_styles*
+            paraphrased variants.  For example, with 2 styles and 2
+            texts: ``[["a_style1", "a_style2"], ["b_style1", "b_style2"]]``.
+        Raises
+        ------
+        ValueError
+            If ``originals`` and ``paraphrases_by_text`` have different
+            lengths.
+        """
+        if len(originals) != len(paraphrases_by_text):
+            raise ValueError(
+                f"originals has {len(originals)} items but "
+                f"paraphrases_by_text has {len(paraphrases_by_text)}."
+            )
+
+        for i, (orig, paras) in enumerate(zip(originals, paraphrases_by_text)):
+            if len(paras) != self._n_styles:
+                _log.warning(
+                    "Expected %d paraphrases for text %d, got %d.",
+                    self._n_styles, i, len(paras),
+                )
+            record = {"original": orig, "paraphrases": paras}
+            if self._output_path is None:
+                self._accumulated.append(record)
+            else:
+                assert self._handle is not None
+                self._handle.write(
+                    json.dumps(record, ensure_ascii=False) + "\n"
+                )
+
+    def finish(self) -> DiversifyOutput:
+        """Close the file handle and return the final result.
+
+        Returns
+        -------
+        list[dict]
+            When ``output_path`` was ``None`` (in-memory mode).  Each dict
+            has keys ``"original"`` and ``"paraphrases"``.
+        Path
+            When results were written to disk — the ``.jsonl`` path.
+        """
+        if self._handle is not None:
+            self._handle.close()
+            self._handle = None
+
+        if self._output_path is None:
+            # In-memory mode: return the accumulated list of dicts.
+            return self._accumulated
+
+        # Disk mode: return the output path.
+        return self._output_path

diversify_text/_postprocess.py

@@ -0,0 +1,64 @@
+"""Text postprocessing utilities for diversify."""
+
+from __future__ import annotations
+
+from diversify_text._preprocess import PreprocessContext
+
+
+def reassemble_segments(
+    segments_per_text: list[list[str]],
+    paraphrases_by_segment: list[list[str]],
+) -> list[list[str]]:
+    """Join per-segment paraphrases back into per-original-text paraphrases.
+
+    Parameters
+    ----------
+    segments_per_text : list[list[str]]
+        The sentence segments for each original text (from
+        :func:`~diversify_text._preprocess.split_sentences`).
+    paraphrases_by_segment : list[list[str]]
+        Flat list of paraphrases for every segment, shape
+        ``[total_segments][n_styles]``.
+
+    Returns
+    -------
+    list[list[str]]
+        Shape ``[n_texts][n_styles]`` — reassembled paraphrases.
+    """
+    result = []
+    seg_idx = 0
+    for segs in segments_per_text:
+        seg_paras = paraphrases_by_segment[seg_idx : seg_idx + len(segs)]
+        n_styles = len(seg_paras[0])
+        result.append([" ".join(sp[i] for sp in seg_paras) for i in range(n_styles)])
+        seg_idx += len(segs)
+    return result
+
+
+def postprocess(
+    candidate: list[list[str]],
+    context: PreprocessContext,
+) -> list[list[str]]:
+    """Undo preprocessing transformations on a candidate set.
+
+    Applies the inverse of each step performed by
+    :func:`~diversify_text._preprocess.preprocess`, using the state stored in
+    *context*.
+
+    Parameters
+    ----------
+    candidate : list[list[str]]
+        Raw generation output, shape ``[n_generation_texts][n_styles]``.
+    context : PreprocessContext
+        Context returned by :func:`~diversify_text._preprocess.preprocess`.
+
+    Returns
+    -------
+    list[list[str]]
+        Shape ``[n_texts][n_styles]`` — one paraphrase per original text
+        per style.
+    """
+    if context.segments_per_text is not None:
+        candidate = reassemble_segments(context.segments_per_text, candidate)
+
+    return candidate

diversify_text/_preprocess.py

@@ -0,0 +1,76 @@
+"""Text preprocessing utilities for diversify."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import pysbd  # https://github.com/nipunsadvilkar/pySBD, published at EMNLP 2020, rule-based
+
+
+_SEGMENTER = pysbd.Segmenter(language="en", clean=False)
+
+
+@dataclass
+class PreprocessContext:
+    """State produced by :func:`preprocess` and consumed by
+    :func:`~diversify_text._postprocess.postprocess`.
+
+    New preprocessing steps can add fields here without changing the
+    caller in ``core.py``.
+    """
+
+    segments_per_text: list[list[str]] | None = None
+
+
+def split_sentences(text: str) -> list[str]:
+    """Split *text* into sentences using pysbd.
+
+    Returns a list of stripped sentence strings.  If the text is empty or
+    whitespace-only, returns a single-element list containing the stripped
+    (possibly empty) input.
+    """
+    segments = _SEGMENTER.segment(text.strip())
+    cleaned = [s.strip() for s in segments if s.strip()]
+    return cleaned or [text.strip()]
+
+
+def preprocess(
+    texts: list[str],
+    *,
+    split_on_punctuation: bool = False,
+) -> tuple[list[str], PreprocessContext]:
+    """Prepare a batch of texts for generation.
+
+    Returns the (possibly transformed) texts to feed into the generation
+    method, together with a :class:`PreprocessContext` that
+    :func:`~diversify_text._postprocess.postprocess` needs to undo the
+    transformations.
+
+    Parameters
+    ----------
+    texts : list[str]
+        Original input texts.
+    split_on_punctuation : bool
+        If ``True``, split each text into sentence-level segments and
+        flatten the result.  The per-text segment mapping is stored in
+        the context so that :func:`~diversify_text._postprocess.postprocess`
+        can reassemble them.
+
+    Returns
+    -------
+    generation_texts : list[str]
+        Texts to pass to the generation method.
+    context : PreprocessContext
+        Context needed by :func:`~diversify_text._postprocess.postprocess`.
+    """
+    context = PreprocessContext()
+
+    if split_on_punctuation:
+        context.segments_per_text = [split_sentences(t) for t in texts]
+        generation_texts = [
+            seg for segs in context.segments_per_text for seg in segs
+        ]
+    else:
+        generation_texts = texts
+
+    return generation_texts, context

diversify_text/_utils.py

@@ -0,0 +1,27 @@
+"""Shared internal utilities for diversify."""
+
+from __future__ import annotations
+
+import contextlib
+import logging
+import warnings
+
+
+@contextlib.contextmanager
+def suppress_hf_load_noise():
+    """Silence harmless noise emitted when loading HuggingFace models.
+
+    Covers two sources that Python's warnings module alone cannot reach:
+
+    - Tied-weights notices from the ``transformers`` logging system.
+    - Unexpected-key load reports from the style-embedding model.
+    """
+    transformers_logger = logging.getLogger("transformers")
+    prev_level = transformers_logger.level
+    transformers_logger.setLevel(logging.ERROR)
+    try:
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", message=".*tie.*weight.*")
+            yield
+    finally:
+        transformers_logger.setLevel(prev_level)