openai-gabriel 1.0.1__py3-none-any.whl

gabriel/tasks/deduplicate.py

@@ -0,0 +1,341 @@
+from __future__ import annotations
+
+import asyncio
+import hashlib
+import os
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+from scipy.cluster.vq import kmeans2
+
+from ..core.prompt_template import PromptTemplate, resolve_template
+from ..utils.openai_utils import get_all_responses
+from ..utils import (
+    safest_json,
+    safe_json,
+    get_all_embeddings,
+    warn_if_modality_mismatch,
+)
+from ..utils.logging import announce_prompt_rendering
+
+
+@dataclass
+class DeduplicateConfig:
+    """Configuration for :class:`Deduplicate`."""
+
+    save_dir: str = "deduplicate"
+    file_name: str = "deduplicate_responses.csv"
+    model: str = "gpt-5-mini"
+    n_parallels: int = 650
+    n_runs: int = 3
+    use_dummy: bool = False
+    max_timeout: Optional[float] = None
+    additional_instructions: Optional[str] = None
+    use_embeddings: bool = True
+    group_size: int = 500
+    modality: str = "entity"
+    max_words_per_text: int = 500
+
+    def __post_init__(self) -> None:
+        if self.additional_instructions is not None:
+            cleaned = str(self.additional_instructions).strip()
+            self.additional_instructions = cleaned or None
+
+
+class Deduplicate:
+    """LLM-assisted deduplication for a single DataFrame column."""
+
+    def __init__(
+        self,
+        cfg: DeduplicateConfig,
+        template: Optional[PromptTemplate] = None,
+        template_path: Optional[str] = None,
+    ) -> None:
+        expanded = Path(os.path.expandvars(os.path.expanduser(cfg.save_dir)))
+        expanded.mkdir(parents=True, exist_ok=True)
+        cfg.save_dir = str(expanded)
+        self.cfg = cfg
+        self.template = resolve_template(
+            template=template,
+            template_path=template_path,
+            reference_filename="deduplicate_prompt.jinja2",
+        )
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _deduplicate(series: pd.Series) -> Tuple[List[str], Dict[str, List[str]], Dict[str, str]]:
+        """Return (unique_values, rep_to_group, orig_to_rep) for a Series."""
+        rep_map: Dict[str, str] = {}
+        groups: Dict[str, List[str]] = {}
+        orig_to_rep: Dict[str, str] = {}
+        for val in series.dropna().astype(str):
+            norm = re.sub(r"[^0-9a-z]+", "", val.lower())
+            if norm in rep_map:
+                rep = rep_map[norm]
+                groups[rep].append(val)
+            else:
+                rep_map[norm] = val
+                groups[val] = [val]
+            orig_to_rep[val] = rep_map[norm]
+        uniques = list(groups.keys())
+        return uniques, groups, orig_to_rep
+
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _print_stats(before: pd.Series, after: pd.Series, *, run_idx: int, total_runs: int) -> None:
+        total = before.notna().sum()
+        diff = (before.fillna("<NA>") != after.fillna("<NA>")).sum()
+        percent = (diff / total * 100) if total else 0.0
+        unique_mapped = after.dropna().nunique()
+        avg_per_map = (total / unique_mapped) if unique_mapped else 0.0
+        print(
+            f"[Deduplicate] Run {run_idx + 1}/{total_runs}: {diff} deduplications "
+            f"({percent:.2f}% of {total})."
+        )
+        print(
+            f"[Deduplicate] Unique mapped terms: {unique_mapped}; "
+            f"avg terms per mapping: {avg_per_map:.2f}."
+        )
+
+    # ------------------------------------------------------------------
+    async def _run_once(
+        self,
+        df_proc: pd.DataFrame,
+        *,
+        column_name: str,
+        output_col: str,
+        raw_texts: Optional[Dict[str, str]] = None,
+        reset_files: bool,
+        run_idx: int,
+        **kwargs: Any,
+    ) -> None:
+        uniques, groups, orig_to_rep = self._deduplicate(df_proc[column_name])
+
+        use_embeddings = self.cfg.use_embeddings and len(uniques) >= self.cfg.group_size
+
+        batches: List[List[str]] = []
+        if use_embeddings:
+            embed_texts = uniques
+            if raw_texts is not None:
+                embed_texts = [raw_texts.get(u, u) for u in uniques]
+            emb = await get_all_embeddings(
+                texts=embed_texts,
+                identifiers=uniques,
+                save_path=os.path.join(self.cfg.save_dir, "deduplicate_embeddings.pkl"),
+                reset_file=reset_files and run_idx == 0,
+                use_dummy=self.cfg.use_dummy,
+                verbose=False,
+            )
+            if emb:
+                arr = np.array([emb[u] for u in uniques], dtype=float)
+                k = max(1, int(np.ceil(len(uniques) / self.cfg.group_size)))
+                _, labels = kmeans2(arr, k, minit="points")
+                clusters: List[List[str]] = [[] for _ in range(k)]
+                for term, lbl in zip(uniques, labels):
+                    clusters[int(lbl)].append(term)
+                current: List[str] = []
+                for cluster in clusters:
+                    for term in cluster:
+                        current.append(term)
+                        if len(current) >= self.cfg.group_size:
+                            batches.append(current)
+                            current = []
+                if current:
+                    batches.append(current)
+        if not batches:
+            sorted_uniques = sorted(uniques, key=lambda x: x.lower())
+            for i in range(0, len(sorted_uniques), self.cfg.group_size):
+                batches.append(sorted_uniques[i : i + self.cfg.group_size])
+
+        prompts: List[str] = []
+        identifiers: List[str] = []
+        announce_prompt_rendering("Deduplicate", len(batches) * max(1, self.cfg.n_runs))
+        for idx, items in enumerate(batches):
+            raw_terms: Any = items
+            if raw_texts is not None:
+                raw_terms = {ident: raw_texts.get(ident, "") for ident in items}
+            prompts.append(
+                self.template.render(
+                    group_id=f"deduplicate_{idx:05d}",
+                    raw_terms=raw_terms,
+                    additional_instructions=self.cfg.additional_instructions or "",
+                    modality=self.cfg.modality,
+                )
+            )
+            identifiers.append(f"deduplicate_{idx:05d}")
+
+        base, ext = os.path.splitext(self.cfg.file_name)
+        if self.cfg.n_runs > 1:
+            response_file = f"{base}_run{run_idx + 1}{ext}"
+        else:
+            response_file = self.cfg.file_name
+        save_path = os.path.join(self.cfg.save_dir, response_file)
+        if prompts:
+            resp_df = await get_all_responses(
+                prompts=prompts,
+                identifiers=identifiers,
+                n_parallels=self.cfg.n_parallels,
+                model=self.cfg.model,
+                save_path=save_path,
+                use_dummy=self.cfg.use_dummy,
+                max_timeout=self.cfg.max_timeout,
+                json_mode=True,
+                reset_files=reset_files,
+                **kwargs,
+            )
+        else:
+            resp_df = pd.DataFrame(columns=["Identifier", "Response"])
+
+        resp_map = dict(zip(resp_df.get("Identifier", []), resp_df.get("Response", [])))
+        parsed = await asyncio.gather(
+            *[safest_json(resp_map.get(i, "")) for i in identifiers]
+        )
+
+        mappings: Dict[str, str] = {}
+        for items, res in zip(batches, parsed):
+            if isinstance(res, str):
+                res = safe_json(res)
+            if isinstance(res, dict):
+                for rep, vals in res.items():
+                    if isinstance(vals, list):
+                        for val in vals:
+                            if isinstance(val, str) and val in items:
+                                mappings[val] = rep
+            elif isinstance(res, list):
+                for row in res:
+                    if isinstance(row, str):
+                        row = safe_json(row)
+                    if isinstance(row, dict):
+                        inp = row.get("input")
+                        mapped = row.get("mapped")
+                        if (
+                            isinstance(inp, str)
+                            and isinstance(mapped, str)
+                            and inp in items
+                        ):
+                            mappings[inp] = mapped
+
+        for rep in uniques:
+            mappings.setdefault(rep, rep)
+
+        mapped_vals: List[Optional[str]] = []
+        for val in df_proc[column_name]:
+            if pd.isna(val):
+                mapped_vals.append(val)
+            else:
+                rep = orig_to_rep.get(str(val), str(val))
+                mapped_vals.append(mappings.get(rep, rep))
+        df_proc[output_col] = mapped_vals
+
+    # ------------------------------------------------------------------
+
+    async def run(
+        self,
+        df: pd.DataFrame,
+        *,
+        column_name: str,
+        reset_files: bool = False,
+        nruns: Optional[int] = None,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        """Deduplicate a column using LLM assistance."""
+
+        df_proc = df.reset_index(drop=True).copy()
+        n_runs = nruns if nruns is not None else self.cfg.n_runs
+        values = df_proc[column_name].tolist()
+        warn_if_modality_mismatch(values, self.cfg.modality, column_name=column_name)
+        current_col = column_name
+        raw_texts: Optional[Dict[str, str]] = None
+        id_to_original: Dict[str, Any] = {}
+
+        if self.cfg.modality == "text":
+            if self.cfg.group_size > 25:
+                print(
+                    "[Deduplicate] For modality='text', consider a smaller group_size "
+                    "(e.g., 25) to keep prompts concise."
+                )
+            ids: List[Optional[str]] = []
+            truncated: List[Optional[str]] = []
+            truncated_count = 0
+            raw_texts = {}
+            for value in values:
+                if pd.isna(value):
+                    ids.append(None)
+                    truncated.append(None)
+                    continue
+                text = str(value)
+                sha8 = hashlib.sha1(text.encode()).hexdigest()[:8]
+                ids.append(sha8)
+                if sha8 not in id_to_original:
+                    id_to_original[sha8] = value
+                words = text.split()
+                if len(words) > self.cfg.max_words_per_text:
+                    truncated_count += 1
+                    clipped = " ".join(words[: self.cfg.max_words_per_text])
+                else:
+                    clipped = text
+                truncated.append(clipped)
+                raw_texts[sha8] = clipped
+            df_proc["_dedup_id"] = ids
+            df_proc[f"{column_name}_truncated"] = truncated
+            total = sum(1 for v in values if not pd.isna(v))
+            frac = (truncated_count / total * 100) if total else 0.0
+            print(
+                f"[Deduplicate] Truncated {truncated_count}/{total} texts "
+                f"({frac:.2f}%) to max_words_per_text={self.cfg.max_words_per_text}."
+            )
+            current_col = "_dedup_id"
+        for i in range(n_runs):
+            if self.cfg.modality == "text":
+                base = f"mapped_{column_name}_ids"
+                if n_runs == 1:
+                    output_col = base
+                elif i == n_runs - 1:
+                    output_col = f"{base}_final"
+                else:
+                    output_col = f"{base}_run{i + 1}"
+            else:
+                if n_runs == 1:
+                    output_col = f"mapped_{column_name}"
+                elif i == n_runs - 1:
+                    output_col = f"mapped_{column_name}_final"
+                else:
+                    output_col = f"mapped_{column_name}_run{i + 1}"
+            await self._run_once(
+                df_proc,
+                column_name=current_col,
+                output_col=output_col,
+                raw_texts=raw_texts,
+                reset_files=reset_files,
+                run_idx=i,
+                **kwargs,
+            )
+            self._print_stats(
+                df_proc[current_col],
+                df_proc[output_col],
+                run_idx=i,
+                total_runs=n_runs,
+            )
+            current_col = output_col
+        if self.cfg.modality == "text":
+            if n_runs > 1:
+                df_proc[f"mapped_{column_name}_ids"] = df_proc[current_col]
+            final_ids = (
+                df_proc[f"mapped_{column_name}_ids"]
+                if f"mapped_{column_name}_ids" in df_proc.columns
+                else df_proc[current_col]
+            )
+            mapped_texts: List[Any] = []
+            for val in final_ids:
+                if pd.isna(val):
+                    mapped_texts.append(val)
+                else:
+                    mapped_texts.append(id_to_original.get(str(val), val))
+            df_proc[f"mapped_{column_name}"] = mapped_texts
+        elif n_runs > 1:
+            df_proc[f"mapped_{column_name}"] = df_proc[current_col]
+        return df_proc

gabriel/tasks/deidentify.py

@@ -0,0 +1,316 @@
+from __future__ import annotations
+
+import ast
+import json
+import os
+import re
+from copy import deepcopy
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple
+
+import pandas as pd
+
+from ..core.prompt_template import PromptTemplate, resolve_template
+from ..utils import safest_json
+from ..utils.openai_utils import get_all_responses
+from ..utils.logging import announce_prompt_rendering
+
+
+# ────────────────────────────
+# Configuration dataclass
+# ────────────────────────────
+@dataclass
+class DeidentifyConfig:
+    """Configuration for :class:`Deidentifier`."""
+
+    model: str = "gpt-5-mini"
+    n_parallels: int = 650
+    save_dir: str = "deidentify"
+    file_name: str = "deidentified.csv"
+    use_dummy: bool = False
+    max_timeout: Optional[float] = None
+    max_words_per_call: int = 7500
+    additional_instructions: Optional[str] = None
+    reasoning_effort: Optional[str] = None
+    reasoning_summary: Optional[str] = None
+    n_passes: int = 1
+    use_existing_mappings_only: bool = False
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.n_passes, int) or self.n_passes < 1:
+            raise ValueError("n_passes must be an integer >= 1")
+        if self.additional_instructions is not None:
+            cleaned = str(self.additional_instructions).strip()
+            self.additional_instructions = cleaned or None
+
+
+# ────────────────────────────
+# Main de-identification task
+# ────────────────────────────
+class Deidentifier:
+    """Iterative de-identification of sensitive entities in text."""
+
+    def __init__(
+        self,
+        cfg: DeidentifyConfig,
+        template: Optional[PromptTemplate] = None,
+        template_path: Optional[str] = None,
+    ) -> None:
+        expanded = Path(os.path.expandvars(os.path.expanduser(cfg.save_dir)))
+        expanded.mkdir(parents=True, exist_ok=True)
+        cfg.save_dir = str(expanded)
+        self.cfg = cfg
+        self.template = resolve_template(
+            template=template,
+            template_path=template_path,
+            reference_filename="deidentification_prompt.jinja2",
+        )
+
+    # ------------------------------------------------------------------
+    # Helpers
+    # ------------------------------------------------------------------
+    @staticmethod
+    def _chunk_by_words(text: str, max_words: int) -> List[str]:
+        words = text.split()
+        if len(words) <= max_words:
+            return [text]
+        return [" ".join(words[i : i + max_words]) for i in range(0, len(words), max_words)]
+
+    @staticmethod
+    def _coerce_mapping(value: Any) -> Optional[Dict[str, Any]]:
+        """Attempt to convert ``value`` into a mapping dictionary."""
+
+        if value is None:
+            return None
+        try:
+            if pd.isna(value):  # type: ignore[arg-type]
+                return None
+        except Exception:
+            pass
+
+        if isinstance(value, dict):
+            return deepcopy(value)
+
+        if isinstance(value, str):
+            cleaned = value.strip()
+            if not cleaned or cleaned.lower() in {"nan", "none"}:
+                return None
+            try:
+                parsed = json.loads(cleaned)
+            except Exception:
+                try:
+                    parsed = ast.literal_eval(cleaned)
+                except Exception:
+                    return None
+            if isinstance(parsed, dict):
+                return deepcopy(parsed)
+            return None
+
+        return None
+
+    # ------------------------------------------------------------------
+    # Main entry point
+    # ------------------------------------------------------------------
+    async def run(
+        self,
+        df: pd.DataFrame,
+        column_name: str,
+        *,
+        grouping_column: Optional[str] = None,
+        mapping_column: Optional[str] = None,
+        reset_files: bool = False,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        """Deidentify all texts in ``df[column_name]``.
+
+        Parameters
+        ----------
+        df:
+            Input DataFrame.
+        column_name:
+            Name of the column containing the text to de-identify.
+        grouping_column:
+            Optional column whose values determine which rows belong to the same
+            individual/entity. When omitted, each row is treated independently.
+        mapping_column:
+            Optional column containing pre-existing mapping dictionaries. The
+            first non-empty mapping encountered for each group is used as the
+            warm start and is also the mapping reused when
+            ``use_existing_mappings_only`` is ``True``.
+        reset_files:
+            When ``True``, intermediate CSV logs from :func:`get_all_responses`
+            are regenerated.
+        **kwargs:
+            Additional keyword arguments forwarded to
+            :func:`gabriel.utils.openai_utils.get_all_responses`.
+        """
+
+        if column_name not in df.columns:
+            raise ValueError(f"Column '{column_name}' not found in DataFrame")
+        if grouping_column is not None and grouping_column not in df.columns:
+            raise ValueError(
+                f"Grouping column '{grouping_column}' not found in DataFrame"
+            )
+        if mapping_column is not None and mapping_column not in df.columns:
+            raise ValueError(
+                f"Mapping column '{mapping_column}' not found in DataFrame"
+            )
+
+        df_proc = df.reset_index(drop=True).copy()
+
+        if grouping_column is None:
+            df_proc["group_id"] = df_proc.index.astype(str)
+        else:
+            df_proc["group_id"] = df_proc[grouping_column].astype(str)
+
+        group_ids = df_proc["group_id"].unique().tolist()
+        base_name = Path(self.cfg.file_name).stem
+        csv_path = Path(self.cfg.save_dir) / f"{base_name}_cleaned.csv"
+        raw_prefix = Path(self.cfg.save_dir) / f"{base_name}_raw_responses"
+
+        group_segments: Dict[str, List[str]] = {}
+        for gid in group_ids:
+            segs: List[str] = []
+            texts = (
+                df_proc.loc[df_proc["group_id"] == gid, column_name]
+                .fillna("")
+                .astype(str)
+                .tolist()
+            )
+            for text in texts:
+                segs.extend(self._chunk_by_words(text, self.cfg.max_words_per_call))
+            group_segments[gid] = segs
+
+        group_to_map: Dict[str, dict] = {gid: {} for gid in group_ids}
+
+        warm_start_count = 0
+        if mapping_column is not None:
+            for gid, subset in df_proc.groupby("group_id"):
+                values = subset[mapping_column].tolist()
+                mapping = next(
+                    (m for m in (self._coerce_mapping(v) for v in values) if m is not None),
+                    None,
+                )
+                if mapping is not None:
+                    group_to_map[str(gid)] = mapping
+                    warm_start_count += 1
+            if warm_start_count:
+                print(
+                    f"[Deidentify] Loaded {warm_start_count} warm-start mapping(s) "
+                    f"from column '{mapping_column}'."
+                )
+            else:
+                print(
+                    f"[Deidentify] Column '{mapping_column}' provided but no usable "
+                    "mappings were found."
+                )
+
+        print(
+            "[Deidentify] Tip: edit the first mapping for each group and rerun with "
+            "use_existing_mappings_only=True to apply your changes without new LLM calls."
+        )
+
+        max_rounds = max(len(segs) for segs in group_segments.values()) if group_segments else 0
+
+        if self.cfg.use_existing_mappings_only:
+            print(
+                "[Deidentify] use_existing_mappings_only=True -> skipping LLM calls and "
+                "reusing provided mappings."
+            )
+            missing = [gid for gid, mapping in group_to_map.items() if not mapping]
+            if missing:
+                print(
+                    "[Deidentify] Warning: no mapping provided for "
+                    f"{len(missing)} group(s). Their text will be returned unchanged."
+                )
+        else:
+            for pass_idx in range(self.cfg.n_passes):
+                if self.cfg.n_passes > 1:
+                    print(
+                        f"[Deidentify] Starting pass {pass_idx + 1}/{self.cfg.n_passes}."
+                    )
+
+                for rnd in range(max_rounds):
+                    prompts: List[str] = []
+                    identifiers: List[str] = []
+                    id_to_gid: Dict[str, str] = {}
+
+                    for gid in group_ids:
+                        segs = group_segments[gid]
+                        if rnd >= len(segs):
+                            continue
+
+                        ident = f"{gid}_pass{pass_idx}_seg_{rnd}"
+                        identifiers.append(ident)
+                        id_to_gid[ident] = gid
+                        prompts.append(
+                            self.template.render(
+                                text=segs[rnd],
+                                current_map=json.dumps(
+                                    group_to_map.get(gid, {}), ensure_ascii=False
+                                ),
+                                additional_instructions=self.cfg.additional_instructions,
+                            )
+                        )
+
+                    if not prompts:
+                        continue
+
+                    announce_prompt_rendering("Deidentify", len(prompts))
+
+                    save_path = raw_prefix.with_name(
+                        f"{raw_prefix.name}_pass{pass_idx}_round{rnd}.csv"
+                    )
+                    batch_df = await get_all_responses(
+                        prompts=prompts,
+                        identifiers=identifiers,
+                        n_parallels=self.cfg.n_parallels,
+                        model=self.cfg.model,
+                        save_path=str(save_path),
+                        use_dummy=self.cfg.use_dummy,
+                        max_timeout=self.cfg.max_timeout,
+                        json_mode=True,
+                        reasoning_effort=self.cfg.reasoning_effort,
+                        reasoning_summary=self.cfg.reasoning_summary,
+                        reset_files=reset_files,
+                        **kwargs,
+                    )
+
+                    for ident, resp in zip(batch_df["Identifier"], batch_df["Response"]):
+                        gid = id_to_gid.get(ident)
+                        if gid is None:
+                            continue
+                        main = resp[0] if isinstance(resp, list) and resp else ""
+                        parsed = await safest_json(main)
+                        if isinstance(parsed, dict):
+                            group_to_map[gid] = parsed
+
+        mappings_col: List[dict] = []
+        deidentified_texts: List[str] = []
+        for _, row in df_proc.iterrows():
+            gid = row["group_id"]
+            mapping = group_to_map.get(gid, {})
+            mappings_col.append(mapping)
+            text = str(row[column_name])
+            deid_text = text
+            pairs: List[Tuple[str, str]] = []
+            for entry in mapping.values():
+                if isinstance(entry, dict):
+                    casted = entry.get("casted form", "")
+                    for real in entry.get("real forms", []) or []:
+                        if casted and real:
+                            pairs.append((real, casted))
+            pairs.sort(key=lambda x: len(x[0]), reverse=True)
+            for real, casted in pairs:
+                escaped = re.escape(real)
+                pattern = re.compile(
+                    rf"(?<!\w){escaped}(?!\w)", flags=re.IGNORECASE
+                )
+                deid_text = pattern.sub(casted, deid_text)
+            deidentified_texts.append(deid_text)
+
+        df_proc["mapping"] = mappings_col
+        df_proc["deidentified_text"] = deidentified_texts
+        df_proc.to_csv(csv_path, index=False)
+        return df_proc