From d5046ee855a5a4ce45a4274974366bda524aa8a7 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Wed, 7 Jan 2026 12:16:38 +0530 Subject: [PATCH 01/14] added embedding for map --- datatune/core/dask/map_dask.py | 31 ++++++++++++++++++++++++++++--- datatune/core/map.py | 3 ++- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/datatune/core/dask/map_dask.py b/datatune/core/dask/map_dask.py index 0d5dfa1..22cc540 100644 --- a/datatune/core/dask/map_dask.py +++ b/datatune/core/dask/map_dask.py @@ -40,6 +40,7 @@ def llm_batch_inference( prompt: str, serialized_input_column: str, expected_new_fields: List[str], + merge_data: List[dict], df: pd.DataFrame, ) -> pd.DataFrame: """ @@ -74,12 +75,33 @@ def llm_batch_inference( "IF A VALUE FOR A COLUMN DOES NOT EXIST SET IT TO None" \ "'index=|{key1: None, key2: value2, ...}'" ) + input_series = df[serialized_input_column] - df[llm_output_column] = llm( - df[serialized_input_column], prefix, prompt, suffix, optimized=True + dup_to_canon = { + dup: c["canonical_id"] + for c in merge_data + for dup in c["duplicate_ids"] + } + + canonical_idx = input_series.index.difference(dup_to_canon.keys()) + + + canonical_input = input_series.loc[canonical_idx] + + llm_out = llm( + canonical_input, + prefix, + prompt, + suffix, + optimized=True ) - return df + df.loc[canonical_idx, llm_output_column] = llm_out + + for dup, canon in dup_to_canon.items(): + df.loc[dup, llm_output_column] = df.loc[canon, llm_output_column] + + return df def parse_llm_output(llm_output: Union[str, Exception]) -> Union[Dict, Exception]: """ @@ -190,11 +212,13 @@ def __init__( input_fields: Optional[List] = None, output_fields: Optional[List] = None, name: Optional[str] = None, + merge_data: List[dict] = None, ): super().__init__(name=name) self.prompt = prompt self.input_fields = input_fields self.output_fields = output_fields + self.merge_data = merge_data or [] self.serialized_input_column = f"{self.name}_SERIALIZED_INPUT__DATATUNE__" self.prompt_column = f"{self.name}_MAP_PROMPT__DATATUNE__" self.llm_output_column = f"{self.name}_LLM_OUTPUT__DATATUNE__" @@ -234,6 +258,7 @@ def __call__(self, llm: Callable, df: Dict): self.prompt, self.serialized_input_column, self.output_fields, + self.merge_data ), meta=meta_dict, ) diff --git a/datatune/core/map.py b/datatune/core/map.py index a5113ef..2449a56 100644 --- a/datatune/core/map.py +++ b/datatune/core/map.py @@ -12,7 +12,7 @@ def _is_dask_df(obj): except ImportError: return False -def map(*, prompt, output_fields, input_fields=None): +def map(*, prompt, output_fields, input_fields=None, merge_data=None): def apply(llm, data): if _is_dask_df(data): @@ -21,6 +21,7 @@ def apply(llm, data): prompt=prompt, output_fields=output_fields, input_fields=input_fields, + merge_data=merge_data )(llm, data) elif _is_ibis_table(data): from .ibis.map_ibis import _map_ibis From 343b52a2c28d28fffef77a0dcea799628fdf218e Mon Sep 17 00:00:00 2001 From: fadil4u Date: Wed, 7 Jan 2026 20:44:05 +0530 Subject: [PATCH 02/14] deduplicator --- datatune/__init__.py | 3 +- datatune/core/deduplication.py | 196 +++++++++++++++++++++++++++++++++ 2 files changed, 198 insertions(+), 1 deletion(-) create mode 100644 datatune/core/deduplication.py diff --git a/datatune/__init__.py b/datatune/__init__.py index e66a30e..2e2b747 100644 --- a/datatune/__init__.py +++ b/datatune/__init__.py @@ -2,4 +2,5 @@ from datatune.core.filter import filter from datatune.core.map import map from datatune.core.dask.op import finalize -__all__ = ["map", "filter", "finalize", "Agent"] +from datatune.core.deduplication import Dedup +__all__ = ["map", "filter", "finalize", "Agent", "Dedup"] \ No newline at end of file diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py new file mode 100644 index 0000000..5ea4a09 --- /dev/null +++ b/datatune/core/deduplication.py @@ -0,0 +1,196 @@ +import faiss +import ast +from litellm import embedding +import numpy as np +from datatune.logger import get_logger +from typing import List, Callable, Optional +from functools import partial +import pandas as pd +import dask.dataframe as dd + +logger = get_logger(__name__) +def input_as_string( + serialized_input_column: str, + df: pd.DataFrame, + input_fields: Optional[List[str]] = None, +) -> pd.DataFrame: + """ + Converts each row in the DataFrame to a string representation and stores it in a new column. + + Args: + serialized_input_column (str): Name of the column to store the serialized row data. + df (pd.DataFrame): Input DataFrame to process. + input_fields (Optional[List], optional): List of input fields to include. Defaults to None. + + Returns: + pd.DataFrame: DataFrame with the added serialized input column. + """ + df[serialized_input_column] = [ + str(row.to_dict()) for _, row in df.iterrows() + ] + return df +class SemanticDeduplicator: + def __init__( + self, + embedding_model: str = "text-embedding-3-small", + sim_threshold: float = 0.90, + top_k: int = 50, + hnsw_m: int = 32, + ef_search: int = 64, +): + self.embedding_model = embedding_model + self.sim_threshold = sim_threshold + self.top_k = top_k + self.hnsw_m = hnsw_m + self.ef_search = ef_search + + def _embed_and_cluster(self, input_rows: List[str]): + dicts = [ast.literal_eval(row) for row in input_rows] + + def dict_to_text(d): + return ", ".join(f"{k}: {v}" for k, v in d.items()) + + texts = [dict_to_text(d) for d in dicts] + + response = embedding( + model=self.embedding_model, + input=texts + ) + + embeddings = [item["embedding"] for item in response["data"]] + X = np.array(embeddings).astype("float32") + faiss.normalize_L2(X) + + dim = X.shape[1] + index = faiss.IndexHNSWFlat( + dim, self.hnsw_m, faiss.METRIC_INNER_PRODUCT + ) + index.add(X) + index.hnsw.efSearch = self.ef_search + + clusters = [] + visited = set() + + for i in range(len(X)): + if i in visited: + continue + + D, I = index.search(X[i:i+1], self.top_k) + + group = [] + for score, j in zip(D[0], I[0]): + if score >= self.sim_threshold: + group.append(j) + visited.add(j) + + clusters.append(group) + + return [c for c in clusters if len(c) > 1] + + def _llm_evaluation(self, clusters, input_rows, llm): + batch_prefix = ( + "You are given multiple independent CLUSTERS of records." + "For EACH cluster: " + "- Consider ONLY the records inside that cluster." + "- Decide whether any records are duplicates." + "- If duplicates exist, choose ONE canonical record." + """- If no duplicates exist, say "NO_DUPLICATES". must be enclosed in double quotes""" + "Output exactly ONE JSON object PER CLUSTER." + "Do not reference other clusters." + ) + + batch_suffix = """For EACH cluster above: + - Produce exactly ONE output. + - Output MUST correspond to the cluster in the same order. + - If duplicates exist, output valid JSON: + { + "canonical_id": , + "duplicate_ids": [, ...] + } + - If no duplicates exist, output exactly: + "NO_DUPLICATES" + + Do not add explanations. + Do not reference other clusters. + """ + + def cluster_to_string(cluster_id, cluster): + lines = [f"CLUSTER {cluster_id} START"] + for idx in cluster: + lines.append(f"ID {idx}: {input_rows[idx]}") + lines.append(f"CLUSTER {cluster_id} END") + return "\n".join(lines) + + llm_inputs = [ + cluster_to_string(i, cluster) + for i, cluster in enumerate(clusters) + ] + + return llm( + llm_inputs, + batch_prefix=batch_prefix, + prompt_per_row="", + batch_suffix=batch_suffix, + max_retries=3, + optimized=True + ) + def _parse_llm_outputs(self, llm_outputs, candidate_clusters): + merges = [] + + for i, output in enumerate(llm_outputs): + output = output.strip() + + if output == "NO_DUPLICATES": + continue + + try: + data = ast.literal_eval(output) + except (ValueError, SyntaxError): + logger.warning("Invalid JSON for cluster %d: %s", i, output) + continue + + canonical_id = data.get("canonical_id") + duplicate_ids = data.get("duplicate_ids") + + if not isinstance(canonical_id, int): + continue + + if not isinstance(duplicate_ids, list) or \ + not all(isinstance(x, int) for x in duplicate_ids): + continue + + cluster_set = set(candidate_clusters[i]) + if canonical_id not in cluster_set or \ + not set(duplicate_ids).issubset(cluster_set): + continue + + merges.append({ + "canonical_id": canonical_id, + "duplicate_ids": duplicate_ids + }) + + return merges + + def deduplicate(self, df:dd.DataFrame, llm: Callable): + df = df.map_partitions( + partial( + input_as_string, + self.serialized_input_column, + input_fields=self.input_fields, + ) + ) + candidate_clusters = self._embed_and_cluster(input_rows) + if not candidate_clusters: + return [] + + llm_outputs = self._llm_evaluation( + candidate_clusters, input_rows, llm + ) + + merges = self._parse_llm_outputs( + llm_outputs, candidate_clusters + ) + + return merges + + From 5c1380d5e78d92e7c7e93286a04c5faa4726bd53 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Wed, 7 Jan 2026 21:35:51 +0530 Subject: [PATCH 03/14] deduplication for map --- datatune/__init__.py | 4 +- datatune/core/deduplication.py | 293 +++++++++++++++++---------------- 2 files changed, 150 insertions(+), 147 deletions(-) diff --git a/datatune/__init__.py b/datatune/__init__.py index 2e2b747..49abc82 100644 --- a/datatune/__init__.py +++ b/datatune/__init__.py @@ -2,5 +2,5 @@ from datatune.core.filter import filter from datatune.core.map import map from datatune.core.dask.op import finalize -from datatune.core.deduplication import Dedup -__all__ = ["map", "filter", "finalize", "Agent", "Dedup"] \ No newline at end of file +from datatune.core.deduplication import SemanticDeduplicator +__all__ = ["map", "filter", "finalize", "Agent", "SemanticDeduplicator"] \ No newline at end of file diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index 5ea4a09..10d31a9 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -12,7 +12,6 @@ def input_as_string( serialized_input_column: str, df: pd.DataFrame, - input_fields: Optional[List[str]] = None, ) -> pd.DataFrame: """ Converts each row in the DataFrame to a string representation and stores it in a new column. @@ -30,6 +29,7 @@ def input_as_string( ] return df class SemanticDeduplicator: + def __init__( self, embedding_model: str = "text-embedding-3-small", @@ -43,154 +43,157 @@ def __init__( self.top_k = top_k self.hnsw_m = hnsw_m self.ef_search = ef_search - - def _embed_and_cluster(self, input_rows: List[str]): - dicts = [ast.literal_eval(row) for row in input_rows] - - def dict_to_text(d): - return ", ".join(f"{k}: {v}" for k, v in d.items()) - texts = [dict_to_text(d) for d in dicts] - - response = embedding( - model=self.embedding_model, - input=texts - ) - - embeddings = [item["embedding"] for item in response["data"]] - X = np.array(embeddings).astype("float32") - faiss.normalize_L2(X) - - dim = X.shape[1] - index = faiss.IndexHNSWFlat( - dim, self.hnsw_m, faiss.METRIC_INNER_PRODUCT - ) - index.add(X) - index.hnsw.efSearch = self.ef_search - - clusters = [] - visited = set() - - for i in range(len(X)): - if i in visited: - continue - - D, I = index.search(X[i:i+1], self.top_k) - - group = [] - for score, j in zip(D[0], I[0]): - if score >= self.sim_threshold: - group.append(j) - visited.add(j) - - clusters.append(group) - - return [c for c in clusters if len(c) > 1] - - def _llm_evaluation(self, clusters, input_rows, llm): - batch_prefix = ( - "You are given multiple independent CLUSTERS of records." - "For EACH cluster: " - "- Consider ONLY the records inside that cluster." - "- Decide whether any records are duplicates." - "- If duplicates exist, choose ONE canonical record." - """- If no duplicates exist, say "NO_DUPLICATES". must be enclosed in double quotes""" - "Output exactly ONE JSON object PER CLUSTER." - "Do not reference other clusters." - ) - - batch_suffix = """For EACH cluster above: - - Produce exactly ONE output. - - Output MUST correspond to the cluster in the same order. - - If duplicates exist, output valid JSON: - { - "canonical_id": , - "duplicate_ids": [, ...] - } - - If no duplicates exist, output exactly: - "NO_DUPLICATES" - - Do not add explanations. - Do not reference other clusters. - """ - - def cluster_to_string(cluster_id, cluster): - lines = [f"CLUSTER {cluster_id} START"] - for idx in cluster: - lines.append(f"ID {idx}: {input_rows[idx]}") - lines.append(f"CLUSTER {cluster_id} END") - return "\n".join(lines) - - llm_inputs = [ - cluster_to_string(i, cluster) - for i, cluster in enumerate(clusters) - ] - - return llm( - llm_inputs, - batch_prefix=batch_prefix, - prompt_per_row="", - batch_suffix=batch_suffix, - max_retries=3, - optimized=True - ) - def _parse_llm_outputs(self, llm_outputs, candidate_clusters): - merges = [] - - for i, output in enumerate(llm_outputs): - output = output.strip() - - if output == "NO_DUPLICATES": - continue - - try: - data = ast.literal_eval(output) - except (ValueError, SyntaxError): - logger.warning("Invalid JSON for cluster %d: %s", i, output) - continue - - canonical_id = data.get("canonical_id") - duplicate_ids = data.get("duplicate_ids") - - if not isinstance(canonical_id, int): - continue - - if not isinstance(duplicate_ids, list) or \ - not all(isinstance(x, int) for x in duplicate_ids): - continue - - cluster_set = set(candidate_clusters[i]) - if canonical_id not in cluster_set or \ - not set(duplicate_ids).issubset(cluster_set): - continue - - merges.append({ - "canonical_id": canonical_id, - "duplicate_ids": duplicate_ids - }) - - return merges - - def deduplicate(self, df:dd.DataFrame, llm: Callable): - df = df.map_partitions( - partial( - input_as_string, - self.serialized_input_column, - input_fields=self.input_fields, - ) + def _embed_and_cluster(self, input_rows: List[str]): + dicts = [ast.literal_eval(row) for row in input_rows] + + def dict_to_text(d): + return ", ".join(f"{k}: {v}" for k, v in d.items()) + + texts = [dict_to_text(d) for d in dicts] + + response = embedding( + model=self.embedding_model, + input=texts + ) + logger.info("Obtained embeddings for %d records", len(texts)) + embeddings = [item["embedding"] for item in response["data"]] + X = np.array(embeddings).astype("float32") + faiss.normalize_L2(X) + + dim = X.shape[1] + index = faiss.IndexHNSWFlat( + dim, self.hnsw_m, faiss.METRIC_INNER_PRODUCT + ) + index.add(X) + index.hnsw.efSearch = self.ef_search + + clusters = [] + visited = set() + + for i in range(len(X)): + if i in visited: + continue + + D, I = index.search(X[i:i+1], self.top_k) + + group = [] + for score, j in zip(D[0], I[0]): + if score >= self.sim_threshold: + group.append(j) + visited.add(j) + + clusters.append(group) + cluster_candidates = [c for c in clusters if len(c) > 1] + print(cluster_candidates) + return cluster_candidates + + def _llm_evaluation(self, clusters, input_rows, llm): + batch_prefix = ( + "You are given multiple independent CLUSTERS of records." + "For EACH cluster: " + "- Consider ONLY the records inside that cluster." + "- Decide whether any records are duplicates." + "- If duplicates exist, choose ONE canonical record." + """- If no duplicates exist, say "NO_DUPLICATES". must be enclosed in double quotes""" + "Output exactly ONE JSON object PER CLUSTER." + "Do not reference other clusters." + ) + + batch_suffix = """For EACH cluster above: +- Produce exactly ONE output. +- Output MUST correspond to the cluster in the same order. +- If duplicates exist, output valid JSON: +{ +"canonical_id": , +"duplicate_ids": [, ...] +} +- If no duplicates exist, output exactly: +"NO_DUPLICATES" + +Do not add explanations. +Do not reference other clusters. +""" + + def cluster_to_string(cluster_id, cluster): + lines = [f"CLUSTER {cluster_id} START"] + for idx in cluster: + lines.append(f"ID {idx}: {input_rows[idx]}") + lines.append(f"CLUSTER {cluster_id} END") + return "\n".join(lines) + + llm_inputs = [ + cluster_to_string(i, cluster) + for i, cluster in enumerate(clusters) + ] + logger.info("Sending %d clusters to LLM for evaluation...", len(llm_inputs)) + llm_outputs = llm( + llm_inputs, + batch_prefix=batch_prefix, + prompt_per_row="", + batch_suffix=batch_suffix, + max_retries=3, + optimized=True + ) + logger.info("LLM evaluation within clusters completed.") + return llm_outputs + def _parse_llm_outputs(self, llm_outputs, candidate_clusters): + merges = [] + + for i, output in enumerate(llm_outputs): + output = output.strip() + + if output == "NO_DUPLICATES": + continue + + try: + data = ast.literal_eval(output) + except (ValueError, SyntaxError): + logger.warning("Invalid JSON for cluster %d: %s", i, output) + continue + + canonical_id = data.get("canonical_id") + duplicate_ids = data.get("duplicate_ids") + + if not isinstance(canonical_id, int): + continue + + if not isinstance(duplicate_ids, list) or \ + not all(isinstance(x, int) for x in duplicate_ids): + continue + + cluster_set = set(candidate_clusters[i]) + if canonical_id not in cluster_set or \ + not set(duplicate_ids).issubset(cluster_set): + continue + + merges.append({ + "canonical_id": canonical_id, + "duplicate_ids": duplicate_ids + }) + + return merges + + def __call__(self, df:dd.DataFrame, llm: Callable): + df = df.map_partitions( + partial( + input_as_string, + "serialized_input_column", ) - candidate_clusters = self._embed_and_cluster(input_rows) - if not candidate_clusters: - return [] + ) + input_rows = df["serialized_input_column"].compute().tolist() + candidate_clusters = self._embed_and_cluster(input_rows) + if not candidate_clusters: + return [] - llm_outputs = self._llm_evaluation( - candidate_clusters, input_rows, llm - ) + llm_outputs = self._llm_evaluation( + candidate_clusters, input_rows, llm + ) - merges = self._parse_llm_outputs( - llm_outputs, candidate_clusters - ) + merges = self._parse_llm_outputs( + llm_outputs, candidate_clusters + ) - return merges + return merges From 0e54865066e8120d40afc69c3204c15985cdb547 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Wed, 7 Jan 2026 22:09:13 +0530 Subject: [PATCH 04/14] fixes --- datatune/core/deduplication.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index 10d31a9..dfd5809 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -86,6 +86,7 @@ def dict_to_text(d): clusters.append(group) cluster_candidates = [c for c in clusters if len(c) > 1] print(cluster_candidates) + #print(len(cluster_candidates[0])) return cluster_candidates def _llm_evaluation(self, clusters, input_rows, llm): @@ -95,7 +96,7 @@ def _llm_evaluation(self, clusters, input_rows, llm): "- Consider ONLY the records inside that cluster." "- Decide whether any records are duplicates." "- If duplicates exist, choose ONE canonical record." - """- If no duplicates exist, say "NO_DUPLICATES". must be enclosed in double quotes""" + """- If no duplicates exist, output a STRING "NO_DUPLICATES" enclosed in double quotes.""" "Output exactly ONE JSON object PER CLUSTER." "Do not reference other clusters." ) @@ -108,8 +109,7 @@ def _llm_evaluation(self, clusters, input_rows, llm): "canonical_id": , "duplicate_ids": [, ...] } -- If no duplicates exist, output exactly: -"NO_DUPLICATES" +- If no duplicates exist, output exactly a string "NO_DUPLICATES" enclosed in double quio. Do not add explanations. Do not reference other clusters. @@ -145,6 +145,8 @@ def _parse_llm_outputs(self, llm_outputs, candidate_clusters): if output == "NO_DUPLICATES": continue + elif isinstance(output, set): + continue try: data = ast.literal_eval(output) From 05a718815217c77b606406ce548345ead435ec68 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Fri, 9 Jan 2026 23:15:51 +0530 Subject: [PATCH 05/14] dedup for filter --- datatune/core/dask/filter_dask.py | 31 +++++++++++++++++++++++++++++-- datatune/core/dask/map_dask.py | 3 ++- datatune/core/deduplication.py | 27 +++++++++++++++------------ datatune/core/filter.py | 3 ++- 4 files changed, 48 insertions(+), 16 deletions(-) diff --git a/datatune/core/dask/filter_dask.py b/datatune/core/dask/filter_dask.py index 1d39629..c7ab827 100644 --- a/datatune/core/dask/filter_dask.py +++ b/datatune/core/dask/filter_dask.py @@ -39,6 +39,7 @@ def llm_batch_inference( llm_output_column: str, prompt: str, serialized_input_column: str, + merge_data: List[dict], df: pd.DataFrame, ) -> pd.DataFrame: """ @@ -68,9 +69,32 @@ def llm_batch_inference( "ALWAYS STICK TO THE FORMAT index=|{key1: value1, key2: value2, ...} with added key called '__filter__' with value either True to KEEP the record or False to REMOVE it.\n" "IF A VALUE FOR A COLUMN DOES NOT EXIST SET IT TO None" ) - df[llm_output_column] = llm( - df[serialized_input_column], prefix, prompt, suffix, optimized=True + input_series = df[serialized_input_column] + + dup_to_canon = { + dup: c["canonical_id"] + for c in merge_data + for dup in c["duplicate_ids"] + } + + canonical_idx = input_series.index.difference(dup_to_canon.keys()) + + + canonical_input = input_series.loc[canonical_idx] + + llm_out = llm( + canonical_input, + prefix, + prompt, + suffix, + optimized=True ) + + df.loc[canonical_idx, llm_output_column] = llm_out + + for dup, canon in dup_to_canon.items(): + df.loc[dup, llm_output_column] = df.loc[canon, llm_output_column] + return df @@ -191,6 +215,7 @@ class _filter_dask(Op): def __init__( self, prompt: str, + merge_data: List[dict] = None, input_fields: Optional[List] = None, name: Optional[str] = None, on_error: str = "keep", @@ -198,6 +223,7 @@ def __init__( super().__init__(name=name) self.prompt = prompt self.input_fields = input_fields + self.merge_data = merge_data or [] self.serialized_input_column = f"{self.name}_SERIALIZED_INPUT__DATATUNE__" self.prompt_column = f"{self.name}_FILTER_PROMPT__DATATUNE__" self.llm_output_column = f"{self.name}_LLM_OUTPUT__DATATUNE__" @@ -240,6 +266,7 @@ def __call__(self, llm: Callable, df: Dict): self.llm_output_column, self.prompt, self.serialized_input_column, + self.merge_data, ), meta=meta_dict, ) diff --git a/datatune/core/dask/map_dask.py b/datatune/core/dask/map_dask.py index 22cc540..9d549f5 100644 --- a/datatune/core/dask/map_dask.py +++ b/datatune/core/dask/map_dask.py @@ -28,6 +28,7 @@ def input_as_string( pd.DataFrame: DataFrame with the added serialized input column. """ df_inputs = df[input_fields] if input_fields else df + df_inputs = df_inputs.where(df_inputs.notna(), None) df[serialized_input_column] = [ str(row.to_dict()) for _, row in df_inputs.iterrows() ] @@ -72,7 +73,7 @@ def llm_batch_inference( "'index=|{key1: value1, key2: value2, ...}' with added keys of expected new fields if any." "ALWAYS START YOUR RESPONSE WITH 'index=|' WHERE IS THE INDEX OF THE ROW." \ - "IF A VALUE FOR A COLUMN DOES NOT EXIST SET IT TO None" \ + "IF A VALUE FOR A COLUMN DOES NOT EXIST SET IT TO null" \ "'index=|{key1: None, key2: value2, ...}'" ) input_series = df[serialized_input_column] diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index dfd5809..56a458f 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -24,6 +24,8 @@ def input_as_string( Returns: pd.DataFrame: DataFrame with the added serialized input column. """ + df = df.astype(object).where(df.notna(), None) + df[serialized_input_column] = [ str(row.to_dict()) for _, row in df.iterrows() ] @@ -45,6 +47,7 @@ def __init__( self.ef_search = ef_search def _embed_and_cluster(self, input_rows: List[str]): + print(input_rows) dicts = [ast.literal_eval(row) for row in input_rows] def dict_to_text(d): @@ -102,18 +105,18 @@ def _llm_evaluation(self, clusters, input_rows, llm): ) batch_suffix = """For EACH cluster above: -- Produce exactly ONE output. -- Output MUST correspond to the cluster in the same order. -- If duplicates exist, output valid JSON: -{ -"canonical_id": , -"duplicate_ids": [, ...] -} -- If no duplicates exist, output exactly a string "NO_DUPLICATES" enclosed in double quio. - -Do not add explanations. -Do not reference other clusters. -""" + - Produce exactly ONE output. + - Output MUST correspond to the cluster in the same order. + - If duplicates exist, output valid JSON: + { + "canonical_id": , + "duplicate_ids": [, ...] + } + - If no duplicates exist, output exactly a string "NO_DUPLICATES" enclosed in double quio. + + Do not add explanations. + Do not reference other clusters. + """ def cluster_to_string(cluster_id, cluster): lines = [f"CLUSTER {cluster_id} START"] diff --git a/datatune/core/filter.py b/datatune/core/filter.py index 2421fd1..84485d6 100644 --- a/datatune/core/filter.py +++ b/datatune/core/filter.py @@ -12,7 +12,7 @@ def _is_dask_df(obj): except ImportError: return False -def filter(*, prompt, input_fields=None): +def filter(*, prompt, input_fields=None, merge_data=None): def apply(llm, data): if _is_dask_df(data): @@ -20,6 +20,7 @@ def apply(llm, data): return _filter_dask( prompt=prompt, input_fields=input_fields, + merge_data=merge_data )(llm, data) elif _is_ibis_table(data): from .ibis.filter_ibis import _filter_ibis From d57d3b6e12495965aa31719b0a5825772d6b0ebb Mon Sep 17 00:00:00 2001 From: fadil4u Date: Sun, 18 Jan 2026 23:54:21 +0530 Subject: [PATCH 06/14] remove computation --- datatune/core/deduplication.py | 333 +++++++++++++++++++++++++++++---- 1 file changed, 292 insertions(+), 41 deletions(-) diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index 56a458f..ebd38b0 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -1,3 +1,4 @@ +import glob import faiss import ast from litellm import embedding @@ -7,6 +8,8 @@ from functools import partial import pandas as pd import dask.dataframe as dd +import dask +import os logger = get_logger(__name__) def input_as_string( @@ -46,30 +49,201 @@ def __init__( self.hnsw_m = hnsw_m self.ef_search = ef_search - def _embed_and_cluster(self, input_rows: List[str]): - print(input_rows) - dicts = [ast.literal_eval(row) for row in input_rows] - def dict_to_text(d): - return ", ".join(f"{k}: {v}" for k, v in d.items()) + def _embed_and_write_partition( + self, + part, # lazy Dask Series + partition_id: int, + output_dir: str, + ): + print("running2") + pdf = part + print(type(pdf)) - texts = [dict_to_text(d) for d in dicts] + if pdf.empty: + print("empty partition") + return 0 - response = embedding( - model=self.embedding_model, - input=texts - ) - logger.info("Obtained embeddings for %d records", len(texts)) - embeddings = [item["embedding"] for item in response["data"]] - X = np.array(embeddings).astype("float32") + emb_path = f"{output_dir}/embeddings_part_{partition_id}.npy" + idx_path = f"{output_dir}/index_part_{partition_id}.npy" + + if os.path.exists(emb_path) and os.path.exists(idx_path): + return 0 + + row_index = pdf.index.to_numpy() + + def safe_parse(row): + try: + return ast.literal_eval(row) + except Exception: + return {} + + dicts = [safe_parse(row) for row in pdf] + + texts = [ + ", ".join(f"{k}: {v}" for k, v in d.items()) + for d in dicts + ] + + embeddings = [] + + for i in range(0, len(texts), 256): + batch = texts[i:i+256] + resp = embedding(model=self.embedding_model, input=batch) + embeddings.extend([item["embedding"] for item in resp["data"]]) + + X = np.asarray(embeddings, dtype="float32") faiss.normalize_L2(X) + np.save( + f"{output_dir}/embeddings_part_{partition_id}.npy", + X + ) + + np.save( + f"{output_dir}/index_part_{partition_id}.npy", + row_index + ) + + logger.info("Written embeddings for partition %d to disk." % partition_id) + + return len(pdf) + + + def embed_column_to_disk(self, df, column, output_dir): + print("running") + os.makedirs(output_dir, exist_ok=True) + + tasks = [] + + for pid in range(df.npartitions): + part = df[column].get_partition(pid) + + task = dask.delayed(self._embed_and_write_partition)( + part, + pid, + output_dir, + ) + tasks.append(task) + + dask.compute(*tasks) + + + + def build_faiss_index( + self, + embedding_dir: str, + dim: int, + hnsw_m: int, + ef_search: int, + ): + index = faiss.IndexHNSWFlat( + dim, + hnsw_m, + faiss.METRIC_INNER_PRODUCT, + ) + index.hnsw.efConstruction = 200 + index.hnsw.efSearch = ef_search + + id_to_row_index = [] + + emb_files = sorted( + glob.glob(os.path.join(embedding_dir, "embeddings_part_*.npy")) + ) + if not emb_files: + raise RuntimeError("No embedding files found") + + for emb_path in emb_files: + pid = emb_path.split("_part_")[-1].split(".")[0] + idx_path = os.path.join( + embedding_dir, f"index_part_{pid}.npy" + ) + + X = np.load(emb_path, mmap_mode="r") + row_idx = np.load(idx_path, mmap_mode="r") + + if X.shape[1] != dim: + raise ValueError("Embedding dimension mismatch") + + index.add(X) + id_to_row_index.extend(row_idx.tolist()) + + return index, np.asarray(id_to_row_index) + + + def stream_cluster( + self, + index, + id_to_row_index: np.ndarray, + embedding_dir: str, + top_k: int, + sim_threshold: float, + ): + visited = set() + clusters = [] + + emb_files = sorted( + glob.glob(os.path.join(embedding_dir, "embeddings_part_*.npy")) + ) + + faiss_id_cursor = 0 # global FAISS ID counter + + for emb_path in emb_files: + X = np.load(emb_path, mmap_mode="r") + + D, I = index.search(X, top_k) + + for local_i in range(len(X)): + anchor_faiss_id = faiss_id_cursor + local_i + anchor_row = id_to_row_index[anchor_faiss_id] + + if anchor_row in visited: + continue + + group = [] + + for score, nbr_faiss_id in zip(D[local_i], I[local_i]): + if score < sim_threshold: + continue + + nbr_row = id_to_row_index[nbr_faiss_id] + + if nbr_row == anchor_row: + continue + + if nbr_row not in visited: + group.append(nbr_row) + visited.add(nbr_row) + + if group: + group.append(anchor_row) + visited.add(anchor_row) + clusters.append(group) + + faiss_id_cursor += len(X) + + return clusters + + + def _cluster(self, X: np.ndarray): + """ + X: np.ndarray [N, dim], ALREADY L2-normalized + Returns: List[List[int]] where each sublist contains indices + of semantically similar rows (high precision). + """ + dim = X.shape[1] + index = faiss.IndexHNSWFlat( dim, self.hnsw_m, faiss.METRIC_INNER_PRODUCT ) - index.add(X) index.hnsw.efSearch = self.ef_search + index.hnsw.efConstruction = 200 + + # X is already normalized + index.add(X) + + D, I = index.search(X, self.top_k) clusters = [] visited = set() @@ -78,28 +252,60 @@ def dict_to_text(d): if i in visited: continue - D, I = index.search(X[i:i+1], self.top_k) - group = [] - for score, j in zip(D[0], I[0]): - if score >= self.sim_threshold: + + for score, j in zip(D[i], I[i]): + if j == i: + continue + if score >= self.sim_threshold and j not in visited: group.append(j) visited.add(j) - clusters.append(group) - cluster_candidates = [c for c in clusters if len(c) > 1] - print(cluster_candidates) - #print(len(cluster_candidates[0])) - return cluster_candidates - - def _llm_evaluation(self, clusters, input_rows, llm): + if group: + group.append(i) + visited.add(i) + clusters.append(group) + + return clusters + + + + # def embed_column_to_disk( + # self, + # df, + # column: str, + # output_dir: str, + # ): + # os.makedirs(output_dir, exist_ok=True) + + # delayed_tasks = [] + + # for pid in range(df.npartitions): + # part = df[column].get_partition(pid) + + # task = dask.delayed(self._embed_and_write_partition)( + # part.compute(), # ONLY this partition + # pid, + # output_dir, + # ) + + # delayed_tasks.append(task) + + # dask.compute(*delayed_tasks) + + def _llm_evaluation(self, clusters, df_column, llm): + """ + clusters: List[List[int]] of row indices + df_column: Dask Series (the serialized_input_column) + llm: your LLM function (unchanged) + """ batch_prefix = ( "You are given multiple independent CLUSTERS of records." "For EACH cluster: " "- Consider ONLY the records inside that cluster." "- Decide whether any records are duplicates." "- If duplicates exist, choose ONE canonical record." - """- If no duplicates exist, output a STRING "NO_DUPLICATES" enclosed in double quotes.""" + """- If no duplicates exist, output a STRING "NO_DUPLICATES" enclosed in double quotes.""" "Output exactly ONE JSON object PER CLUSTER." "Do not reference other clusters." ) @@ -112,23 +318,49 @@ def _llm_evaluation(self, clusters, input_rows, llm): "canonical_id": , "duplicate_ids": [, ...] } - - If no duplicates exist, output exactly a string "NO_DUPLICATES" enclosed in double quio. + - If no duplicates exist, output exactly a string "NO_DUPLICATES" enclosed in double quotes. Do not add explanations. Do not reference other clusters. """ - def cluster_to_string(cluster_id, cluster): + # def cluster_to_string(cluster_id, cluster): + # # Fetch only the rows needed for this cluster + # input_rows = df_column.loc[cluster].compute().tolist() + # lines = [f"CLUSTER {cluster_id} START"] + # for idx, row in zip(cluster, input_rows): + # lines.append(f"ID {idx}: {row}") + # lines.append(f"CLUSTER {cluster_id} END") + # return "\n".join(lines) + + # Build a mapping: partition -> cluster rows + from collections import defaultdict + all_needed_indices = set(idx for cluster in clusters for idx in cluster) + cluster_rows = defaultdict(list) # cluster_id -> list of (row_idx, text) + + # ---------- scan partition by partition ---------- + for pid in range(df_column.npartitions): + part = df_column.get_partition(pid).compute() + # only keep rows that are actually needed + needed_in_partition = all_needed_indices.intersection(part.index) + for row_idx in needed_in_partition: + for cluster_id, cluster in enumerate(clusters): + if row_idx in cluster: + cluster_rows[cluster_id].append((row_idx, part.loc[row_idx])) + break # each row belongs to exactly one cluster + + # ---------- build cluster strings in correct order ---------- + llm_inputs = [] + for cluster_id, cluster in enumerate(clusters): lines = [f"CLUSTER {cluster_id} START"] + row_map = dict(cluster_rows[cluster_id]) for idx in cluster: - lines.append(f"ID {idx}: {input_rows[idx]}") + if idx in row_map: + lines.append(f"ID {idx}: {row_map[idx]}") lines.append(f"CLUSTER {cluster_id} END") - return "\n".join(lines) + llm_inputs.append("\n".join(lines)) - llm_inputs = [ - cluster_to_string(i, cluster) - for i, cluster in enumerate(clusters) - ] + # ---------- send to LLM ---------- logger.info("Sending %d clusters to LLM for evaluation...", len(llm_inputs)) llm_outputs = llm( llm_inputs, @@ -140,6 +372,7 @@ def cluster_to_string(cluster_id, cluster): ) logger.info("LLM evaluation within clusters completed.") return llm_outputs + def _parse_llm_outputs(self, llm_outputs, candidate_clusters): merges = [] @@ -180,23 +413,41 @@ def _parse_llm_outputs(self, llm_outputs, candidate_clusters): return merges def __call__(self, df:dd.DataFrame, llm: Callable): + df = df.reset_index(drop=True) + df = df.map_partitions( partial( input_as_string, "serialized_input_column", ) ) - input_rows = df["serialized_input_column"].compute().tolist() - candidate_clusters = self._embed_and_cluster(input_rows) - if not candidate_clusters: - return [] - + #input_rows = df["serialized_input_column"].compute().tolist() + self.embed_column_to_disk( + df, + "serialized_input_column", + "embeddings_output_dir" + ) + index, id_to_row_index = self.build_faiss_index( + embedding_dir="embeddings_output_dir", + dim=1536, + hnsw_m=32, + ef_search=128 + ) + clusters = self.stream_cluster( + index=index, + id_to_row_index=id_to_row_index, + embedding_dir="embeddings_output_dir", + top_k=20, + sim_threshold=0.92, + ) + print("Found clusters:", clusters) + print(len(clusters)) llm_outputs = self._llm_evaluation( - candidate_clusters, input_rows, llm + clusters, df["serialized_input_column"], llm ) merges = self._parse_llm_outputs( - llm_outputs, candidate_clusters + llm_outputs, clusters ) return merges From f00aa2020c053afa5a735a78063055e3153525bb Mon Sep 17 00:00:00 2001 From: fadil4u Date: Mon, 19 Jan 2026 00:10:38 +0530 Subject: [PATCH 07/14] dependecy --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 3847728..d905d74 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -18,6 +18,7 @@ dependencies = [ "litellm", "dask[dataframe]", "pandas" + "faiss-cpu", ] [project.optional-dependencies] From 53fde542bab0dc2d0571db2500a94da1297fd18b Mon Sep 17 00:00:00 2001 From: fadil4u Date: Mon, 19 Jan 2026 00:13:40 +0530 Subject: [PATCH 08/14] bug fix --- pyproject.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d905d74..f8d1ee2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,8 +17,8 @@ dependencies = [ "numpy", "litellm", "dask[dataframe]", - "pandas" - "faiss-cpu", + "pandas", + "faiss-cpu" ] [project.optional-dependencies] From 44906ef81ddec30fef1e2ee3d1428d0fe6932249 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Thu, 22 Jan 2026 00:11:29 +0530 Subject: [PATCH 09/14] reduce --- datatune/__init__.py | 4 +- datatune/core/deduplication.py | 98 +++------------------------------- datatune/core/reduce.py | 22 ++++++++ 3 files changed, 30 insertions(+), 94 deletions(-) create mode 100644 datatune/core/reduce.py diff --git a/datatune/__init__.py b/datatune/__init__.py index 49abc82..16b0938 100644 --- a/datatune/__init__.py +++ b/datatune/__init__.py @@ -2,5 +2,5 @@ from datatune.core.filter import filter from datatune.core.map import map from datatune.core.dask.op import finalize -from datatune.core.deduplication import SemanticDeduplicator -__all__ = ["map", "filter", "finalize", "Agent", "SemanticDeduplicator"] \ No newline at end of file +from datatune.core.reduce import reduce +__all__ = ["map", "filter", "finalize", "Agent", "reduce"] \ No newline at end of file diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index ebd38b0..08f9b4d 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -10,6 +10,7 @@ import dask.dataframe as dd import dask import os +from .registry import register_action logger = get_logger(__name__) def input_as_string( @@ -33,8 +34,9 @@ def input_as_string( str(row.to_dict()) for _, row in df.iterrows() ] return df -class SemanticDeduplicator: +@register_action("dedup") +class SemanticDeduplicator(): def __init__( self, embedding_model: str = "text-embedding-3-small", @@ -186,7 +188,7 @@ def stream_cluster( glob.glob(os.path.join(embedding_dir, "embeddings_part_*.npy")) ) - faiss_id_cursor = 0 # global FAISS ID counter + faiss_id_cursor = 0 for emb_path in emb_files: X = np.load(emb_path, mmap_mode="r") @@ -224,75 +226,6 @@ def stream_cluster( return clusters - - def _cluster(self, X: np.ndarray): - """ - X: np.ndarray [N, dim], ALREADY L2-normalized - Returns: List[List[int]] where each sublist contains indices - of semantically similar rows (high precision). - """ - - dim = X.shape[1] - - index = faiss.IndexHNSWFlat( - dim, self.hnsw_m, faiss.METRIC_INNER_PRODUCT - ) - index.hnsw.efSearch = self.ef_search - index.hnsw.efConstruction = 200 - - # X is already normalized - index.add(X) - - D, I = index.search(X, self.top_k) - - clusters = [] - visited = set() - - for i in range(len(X)): - if i in visited: - continue - - group = [] - - for score, j in zip(D[i], I[i]): - if j == i: - continue - if score >= self.sim_threshold and j not in visited: - group.append(j) - visited.add(j) - - if group: - group.append(i) - visited.add(i) - clusters.append(group) - - return clusters - - - - # def embed_column_to_disk( - # self, - # df, - # column: str, - # output_dir: str, - # ): - # os.makedirs(output_dir, exist_ok=True) - - # delayed_tasks = [] - - # for pid in range(df.npartitions): - # part = df[column].get_partition(pid) - - # task = dask.delayed(self._embed_and_write_partition)( - # part.compute(), # ONLY this partition - # pid, - # output_dir, - # ) - - # delayed_tasks.append(task) - - # dask.compute(*delayed_tasks) - def _llm_evaluation(self, clusters, df_column, llm): """ clusters: List[List[int]] of row indices @@ -323,33 +256,17 @@ def _llm_evaluation(self, clusters, df_column, llm): Do not add explanations. Do not reference other clusters. """ - - # def cluster_to_string(cluster_id, cluster): - # # Fetch only the rows needed for this cluster - # input_rows = df_column.loc[cluster].compute().tolist() - # lines = [f"CLUSTER {cluster_id} START"] - # for idx, row in zip(cluster, input_rows): - # lines.append(f"ID {idx}: {row}") - # lines.append(f"CLUSTER {cluster_id} END") - # return "\n".join(lines) - - # Build a mapping: partition -> cluster rows from collections import defaultdict all_needed_indices = set(idx for cluster in clusters for idx in cluster) - cluster_rows = defaultdict(list) # cluster_id -> list of (row_idx, text) - - # ---------- scan partition by partition ---------- + cluster_rows = defaultdict(list) for pid in range(df_column.npartitions): part = df_column.get_partition(pid).compute() - # only keep rows that are actually needed needed_in_partition = all_needed_indices.intersection(part.index) for row_idx in needed_in_partition: for cluster_id, cluster in enumerate(clusters): if row_idx in cluster: cluster_rows[cluster_id].append((row_idx, part.loc[row_idx])) - break # each row belongs to exactly one cluster - - # ---------- build cluster strings in correct order ---------- + break llm_inputs = [] for cluster_id, cluster in enumerate(clusters): lines = [f"CLUSTER {cluster_id} START"] @@ -360,7 +277,6 @@ def _llm_evaluation(self, clusters, df_column, llm): lines.append(f"CLUSTER {cluster_id} END") llm_inputs.append("\n".join(lines)) - # ---------- send to LLM ---------- logger.info("Sending %d clusters to LLM for evaluation...", len(llm_inputs)) llm_outputs = llm( llm_inputs, @@ -440,8 +356,6 @@ def __call__(self, df:dd.DataFrame, llm: Callable): top_k=20, sim_threshold=0.92, ) - print("Found clusters:", clusters) - print(len(clusters)) llm_outputs = self._llm_evaluation( clusters, df["serialized_input_column"], llm ) diff --git a/datatune/core/reduce.py b/datatune/core/reduce.py new file mode 100644 index 0000000..76c6f91 --- /dev/null +++ b/datatune/core/reduce.py @@ -0,0 +1,22 @@ +import dask.dataframe as dd +from .registry import get_action + +_ACTIONS = {} + +def register_action(name): + def decorator(cls): + _ACTIONS[name] = cls + return cls + return decorator + +def get_action(name): + try: + return _ACTIONS[name] + except KeyError: + raise ValueError(f"Unknown action: {name}") + + +def reduce(df, action: str, **kwargs): + cls = get_action(action) + reducer = cls(**kwargs) + return reducer(df) \ No newline at end of file From 74a6b16d005e1ea58e21668b6398d1555d69363c Mon Sep 17 00:00:00 2001 From: fadil4u Date: Thu, 22 Jan 2026 00:27:58 +0530 Subject: [PATCH 10/14] fixes --- datatune/__init__.py | 1 + datatune/core/deduplication.py | 13 ++++++++----- datatune/core/reduce.py | 4 ++-- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/datatune/__init__.py b/datatune/__init__.py index 16b0938..0f288b8 100644 --- a/datatune/__init__.py +++ b/datatune/__init__.py @@ -2,5 +2,6 @@ from datatune.core.filter import filter from datatune.core.map import map from datatune.core.dask.op import finalize +from datatune.core.deduplication import SemanticDeduplicator from datatune.core.reduce import reduce __all__ = ["map", "filter", "finalize", "Agent", "reduce"] \ No newline at end of file diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index 08f9b4d..0e2ab20 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -10,7 +10,7 @@ import dask.dataframe as dd import dask import os -from .registry import register_action +from .reduce import register_action logger = get_logger(__name__) def input_as_string( @@ -39,17 +39,20 @@ def input_as_string( class SemanticDeduplicator(): def __init__( self, + llm:Callable, embedding_model: str = "text-embedding-3-small", sim_threshold: float = 0.90, top_k: int = 50, hnsw_m: int = 32, ef_search: int = 64, + ): self.embedding_model = embedding_model self.sim_threshold = sim_threshold self.top_k = top_k self.hnsw_m = hnsw_m self.ef_search = ef_search + self.llm = llm def _embed_and_write_partition( @@ -226,7 +229,7 @@ def stream_cluster( return clusters - def _llm_evaluation(self, clusters, df_column, llm): + def _llm_evaluation(self, clusters, df_column): """ clusters: List[List[int]] of row indices df_column: Dask Series (the serialized_input_column) @@ -278,7 +281,7 @@ def _llm_evaluation(self, clusters, df_column, llm): llm_inputs.append("\n".join(lines)) logger.info("Sending %d clusters to LLM for evaluation...", len(llm_inputs)) - llm_outputs = llm( + llm_outputs = self.llm( llm_inputs, batch_prefix=batch_prefix, prompt_per_row="", @@ -328,7 +331,7 @@ def _parse_llm_outputs(self, llm_outputs, candidate_clusters): return merges - def __call__(self, df:dd.DataFrame, llm: Callable): + def __call__(self, df:dd.DataFrame): df = df.reset_index(drop=True) df = df.map_partitions( @@ -357,7 +360,7 @@ def __call__(self, df:dd.DataFrame, llm: Callable): sim_threshold=0.92, ) llm_outputs = self._llm_evaluation( - clusters, df["serialized_input_column"], llm + clusters, df["serialized_input_column"] ) merges = self._parse_llm_outputs( diff --git a/datatune/core/reduce.py b/datatune/core/reduce.py index 76c6f91..3879279 100644 --- a/datatune/core/reduce.py +++ b/datatune/core/reduce.py @@ -1,5 +1,5 @@ import dask.dataframe as dd -from .registry import get_action + _ACTIONS = {} @@ -16,7 +16,7 @@ def get_action(name): raise ValueError(f"Unknown action: {name}") -def reduce(df, action: str, **kwargs): +def reduce(df,*, action: str, **kwargs): cls = get_action(action) reducer = cls(**kwargs) return reducer(df) \ No newline at end of file From ddf2fc3f3fd432a5068d7be30c727337ab161368 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Thu, 22 Jan 2026 00:42:08 +0530 Subject: [PATCH 11/14] clusters var --- datatune/core/filter.py | 4 ++-- datatune/core/map.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/datatune/core/filter.py b/datatune/core/filter.py index 84485d6..d60ef1c 100644 --- a/datatune/core/filter.py +++ b/datatune/core/filter.py @@ -12,7 +12,7 @@ def _is_dask_df(obj): except ImportError: return False -def filter(*, prompt, input_fields=None, merge_data=None): +def filter(*, prompt, input_fields=None, clusters=None): def apply(llm, data): if _is_dask_df(data): @@ -20,7 +20,7 @@ def apply(llm, data): return _filter_dask( prompt=prompt, input_fields=input_fields, - merge_data=merge_data + clusters=clusters )(llm, data) elif _is_ibis_table(data): from .ibis.filter_ibis import _filter_ibis diff --git a/datatune/core/map.py b/datatune/core/map.py index 2449a56..75aede7 100644 --- a/datatune/core/map.py +++ b/datatune/core/map.py @@ -12,7 +12,7 @@ def _is_dask_df(obj): except ImportError: return False -def map(*, prompt, output_fields, input_fields=None, merge_data=None): +def map(*, prompt, output_fields, input_fields=None, clusters=None): def apply(llm, data): if _is_dask_df(data): @@ -21,7 +21,7 @@ def apply(llm, data): prompt=prompt, output_fields=output_fields, input_fields=input_fields, - merge_data=merge_data + clusters=clusters )(llm, data) elif _is_ibis_table(data): from .ibis.map_ibis import _map_ibis From 5a7b7964efbacf2dbde99a1a8c98ae52dfb3ba56 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Sun, 25 Jan 2026 14:17:15 +0530 Subject: [PATCH 12/14] docs --- docs/source/Reduce.md | 124 ++++++++++++++++++++++++ examples/deduplication.ipynb | 181 +++++++++++++++++++++++++++++++++++ 2 files changed, 305 insertions(+) create mode 100644 docs/source/Reduce.md create mode 100644 examples/deduplication.ipynb diff --git a/docs/source/Reduce.md b/docs/source/Reduce.md new file mode 100644 index 0000000..267b8e8 --- /dev/null +++ b/docs/source/Reduce.md @@ -0,0 +1,124 @@ +# Reduce + +The `Reduce` operation in Datatune reduces the size of a dataset by grouping, deduplicating, or otherwise collapsing rows based on a specified action. +It is typically used to minimize downstream processing cost (e.g. LLM calls) by identifying canonical rows and eliminating redundant ones. + + +## Basic Usage + +```python +import datatune as dt +from datatune.llm.llm import LLM +import dask.dataframe as dd + +# Initialize LLM +llm = LLM(model_name="openai/gpt-3.5-turbo") + +# Load data +df = dd.read_csv("data.csv") + +# Get deduplication clusters +clusters = dt.reduce(df, action="dedup", embedding_model="text-embedding-3-small", llm=llm) + + +# Apply transformation +mapped_df = dt.map( + prompt="Extract country and city from the address field", + output_fields=["country", "city"] + clusters=clusters # pass clusters +)(llm, df) + +# Process results +mapped_df.compute() +``` +## Parameters +**df** (DataFrame, required): Input dataset to be reduced. + +**action** (str, required): Name of the reduction action to apply. Examples: "dedup" + +**kwargs** (optional): Configuration parameters forwarded to the selected reduction action.The accepted parameters depend on the specific action being used. + +## Actions + +Reduction behavior is defined by actions, which are registered internally and selected by name at runtime. +Reduce accepts generic parameters, plus action-specific parameters +that are interpreted only by the selected action. + +Each action: + +- Is implemented as a callable class + +- Receives configuration via its constructor + +Operates on the input DataFrame when called + + +### dedup + +- Identifies semantically duplicate rows and selects a canonical representative for each group. Outputs a "dedup map" +```python +[{'canonical_id': 5, 'duplicate_ids': [1, 2, 3, 4, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19]}] +``` +This map can be passed Datatune operations which will only send the canonical rows to LLM API for transformation and transmit it's results to the respective duplicate rows thereby reducing tokens and therefore cost. + +The dedup action works by first embedding all input rows into vector representations, which are streamed to disk to avoid loading the full dataset into memory. These embeddings are then indexed using FAISS and searched to identify clusters of semantically similar rows based on the configured similarity threshold. Finally, within each candidate cluster, an LLM evaluation step is applied to confirm true duplicates and select a canonical representative for each cluster. + +#### Dedup Parameters + +When `action="dedup"`, the following additional parameters may be passed +to `reduce`: + +**llm** (Callable, required) + +LLM callable used for semantic evaluation and final duplicate confirmation. + +--- + +**embedding_model** (str, optional) + +Embedding model used to convert rows into vector representations. + +Default: `"text-embedding-3-small"` + +--- + +**sim_threshold** (float, optional) + +Cosine similarity threshold above which two rows are considered potential +duplicates. + +Higher values result in stricter deduplication. + +Default: `0.90` + +--- + +**top_k** (int, optional) + +Number of nearest neighbors retrieved per row during similarity search. + +Default: `50` + +--- + +**hnsw_m** (int, optional) + +HNSW graph connectivity parameter controlling the number of bi-directional +links created for each node. + +Higher values improve recall but increase memory usage. + +Default: `32` + +--- + +**ef_search** (int, optional) + +Size of the dynamic candidate list during HNSW search. + +Higher values improve search accuracy at the cost of latency. + +Default: `64` + +--- + diff --git a/examples/deduplication.ipynb b/examples/deduplication.ipynb new file mode 100644 index 0000000..5487c79 --- /dev/null +++ b/examples/deduplication.ipynb @@ -0,0 +1,181 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "9d96f064", + "metadata": {}, + "source": [ + "# Datatune with Semantic Deduplication" + ] + }, + { + "cell_type": "markdown", + "id": "e16bfaa5", + "metadata": {}, + "source": [ + "Let's start by installing dependencies. We will be using duckdb as our database backend and OpenAI LLM API." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4f08aea7", + "metadata": {}, + "outputs": [], + "source": [ + "!pip install datatune" + ] + }, + { + "cell_type": "markdown", + "id": "6b326c27", + "metadata": {}, + "source": [ + "## Import required libraries" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "5169a466", + "metadata": {}, + "outputs": [], + "source": [ + "import datatune as dt\n", + "import seaborn as sns\n", + "from datatune.llm.llm import OpenAI\n", + "import dask.dataframe as dd\n", + "import os" + ] + }, + { + "cell_type": "markdown", + "id": "412d36ed", + "metadata": {}, + "source": [ + "## Initialize your LLM" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "26156d6d", + "metadata": {}, + "outputs": [], + "source": [ + "os.environ[\"OPENAI_API_KEY\"] = \"your_openai_api_key_here\" # Replace with your actual OpenAI API key\n", + "llm = OpenAI(model=\"gpt-3.5-turbo\", rpm=500, tpm=150000) # Initialize the LLM with your rate limits " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "bddba029", + "metadata": {}, + "outputs": [], + "source": [ + "df = sns.load_dataset(\"flights\") # Load the flights dataset\n", + "df = dd.from_pandas(df)" + ] + }, + { + "cell_type": "markdown", + "id": "1fda49b6", + "metadata": {}, + "source": [ + "## Get deduplication clusters" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "41090e88", + "metadata": {}, + "outputs": [], + "source": [ + "clusters = dt.reduce(df, action=\"dedup\", embedding_model=\"text-embedding-3-small\", llm=llm)" + ] + }, + { + "cell_type": "markdown", + "id": "d9d65000", + "metadata": {}, + "source": [ + "Reduce gets a deduplication map that can be passed to map and filter which sends only canonical rows to the LLM API for transformation and transmits their result to the duplicate rows thereby reducing tokens and therfore cost." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "a9c327ae", + "metadata": {}, + "outputs": [], + "source": [ + "# Create mapping and filtering prompt to transform the dataset\n", + "\n", + "mapping_prompt = \"Add a column passenger_trend_comment that describes the trend in passenger numbers..make the comment descriptive and varied\"\n", + "\n", + "filtering_prompt = \"Based on the passenger_trend_comment column, filter the dataset to include only those months where there is a significant change in passenger numbers compared to the previous month.\"\n" + ] + }, + { + "cell_type": "markdown", + "id": "e6a4ee14", + "metadata": {}, + "source": [ + "## Transforming our dataset" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "fb992c44", + "metadata": {}, + "outputs": [], + "source": [ + "mapped = dt.map(\n", + " prompt = mapping_prompt,\n", + " output_fields=[\"passenger_trend_comment\"], # input fields to be used for mapping\n", + " input_fields=[\"passengers\",\"month\",\"year\"],\n", + " clusters=clusters # pass deduplication clusters\n", + ")(llm, df)\n", + "\n", + "# Now pass the mapped Ibis table expression to filter\n", + "filtered = dt.filter(\n", + " prompt = filtering_prompt,\n", + " input_fields=[\"passenger_trend_comment\"],\n", + " clusters=clusters # pass deduplication clusters\n", + ")(llm, mapped)\n", + "\n", + "result = filtered.execute() # Result is a pandas DataFrame" + ] + }, + { + "cell_type": "markdown", + "id": "f45f59d3", + "metadata": {}, + "source": [ + "## Convert the transformed dataset into CSV" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "52cd8e16", + "metadata": {}, + "outputs": [], + "source": [ + "result.to_csv(\"duckdb_transformed.csv\")\n", + "\n", + "print(result.head())" + ] + } + ], + "metadata": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} From ea4fd9ae6f9259b2e26eca6c34cc951e79850211 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Tue, 27 Jan 2026 23:13:01 +0530 Subject: [PATCH 13/14] reduce return df --- datatune/core/dask/filter_dask.py | 10 +++++----- datatune/core/dask/map_dask.py | 10 +++++----- datatune/core/deduplication.py | 27 ++++++++++++++++++++++----- 3 files changed, 32 insertions(+), 15 deletions(-) diff --git a/datatune/core/dask/filter_dask.py b/datatune/core/dask/filter_dask.py index c7ab827..56b0255 100644 --- a/datatune/core/dask/filter_dask.py +++ b/datatune/core/dask/filter_dask.py @@ -39,7 +39,7 @@ def llm_batch_inference( llm_output_column: str, prompt: str, serialized_input_column: str, - merge_data: List[dict], + clusters: List[dict], df: pd.DataFrame, ) -> pd.DataFrame: """ @@ -73,7 +73,7 @@ def llm_batch_inference( dup_to_canon = { dup: c["canonical_id"] - for c in merge_data + for c in clusters for dup in c["duplicate_ids"] } @@ -215,7 +215,7 @@ class _filter_dask(Op): def __init__( self, prompt: str, - merge_data: List[dict] = None, + clusters: List[dict] = None, input_fields: Optional[List] = None, name: Optional[str] = None, on_error: str = "keep", @@ -223,7 +223,7 @@ def __init__( super().__init__(name=name) self.prompt = prompt self.input_fields = input_fields - self.merge_data = merge_data or [] + self.clusters = clusters or [] self.serialized_input_column = f"{self.name}_SERIALIZED_INPUT__DATATUNE__" self.prompt_column = f"{self.name}_FILTER_PROMPT__DATATUNE__" self.llm_output_column = f"{self.name}_LLM_OUTPUT__DATATUNE__" @@ -266,7 +266,7 @@ def __call__(self, llm: Callable, df: Dict): self.llm_output_column, self.prompt, self.serialized_input_column, - self.merge_data, + self.clusters, ), meta=meta_dict, ) diff --git a/datatune/core/dask/map_dask.py b/datatune/core/dask/map_dask.py index 9d549f5..deb0a43 100644 --- a/datatune/core/dask/map_dask.py +++ b/datatune/core/dask/map_dask.py @@ -41,7 +41,7 @@ def llm_batch_inference( prompt: str, serialized_input_column: str, expected_new_fields: List[str], - merge_data: List[dict], + clusters: List[dict], df: pd.DataFrame, ) -> pd.DataFrame: """ @@ -80,7 +80,7 @@ def llm_batch_inference( dup_to_canon = { dup: c["canonical_id"] - for c in merge_data + for c in clusters for dup in c["duplicate_ids"] } @@ -213,13 +213,13 @@ def __init__( input_fields: Optional[List] = None, output_fields: Optional[List] = None, name: Optional[str] = None, - merge_data: List[dict] = None, + clusters: List[dict] = None, ): super().__init__(name=name) self.prompt = prompt self.input_fields = input_fields self.output_fields = output_fields - self.merge_data = merge_data or [] + self.clusters = clusters or [] self.serialized_input_column = f"{self.name}_SERIALIZED_INPUT__DATATUNE__" self.prompt_column = f"{self.name}_MAP_PROMPT__DATATUNE__" self.llm_output_column = f"{self.name}_LLM_OUTPUT__DATATUNE__" @@ -259,7 +259,7 @@ def __call__(self, llm: Callable, df: Dict): self.prompt, self.serialized_input_column, self.output_fields, - self.merge_data + self.clusters ), meta=meta_dict, ) diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py index 0e2ab20..f123a1f 100644 --- a/datatune/core/deduplication.py +++ b/datatune/core/deduplication.py @@ -45,6 +45,7 @@ def __init__( top_k: int = 50, hnsw_m: int = 32, ef_search: int = 64, + return_df: bool = False, ): self.embedding_model = embedding_model @@ -53,6 +54,7 @@ def __init__( self.hnsw_m = hnsw_m self.ef_search = ef_search self.llm = llm + self.return_df = return_df def _embed_and_write_partition( @@ -61,9 +63,7 @@ def _embed_and_write_partition( partition_id: int, output_dir: str, ): - print("running2") pdf = part - print(type(pdf)) if pdf.empty: print("empty partition") @@ -116,7 +116,6 @@ def safe_parse(row): def embed_column_to_disk(self, df, column, output_dir): - print("running") os.makedirs(output_dir, exist_ok=True) tasks = [] @@ -330,7 +329,20 @@ def _parse_llm_outputs(self, llm_outputs, candidate_clusters): }) return merges - + def _dedup_df(self, df, clusters): + + # build drop-set ONCE + to_drop = set() + for c in clusters: + to_drop.update(c["duplicate_ids"]) + + def _filter(pdf): + return pdf.loc[~pdf.index.isin(to_drop)] + + return df.map_partitions( + _filter, + meta=df._meta + ) def __call__(self, df:dd.DataFrame): df = df.reset_index(drop=True) @@ -366,7 +378,12 @@ def __call__(self, df:dd.DataFrame): merges = self._parse_llm_outputs( llm_outputs, clusters ) - + if self.return_df: + + deduped_df = self._dedup_df(df, merges) + deduped_df = deduped_df.drop(columns=["serialized_input_column"]) + return deduped_df + return merges From 00a9df372326841c0e3fd39e2d94c68e1365f0a2 Mon Sep 17 00:00:00 2001 From: fadil4u Date: Tue, 27 Jan 2026 23:20:36 +0530 Subject: [PATCH 14/14] docs --- docs/source/Reduce.md | 5 +++++ examples/deduplication.ipynb | 21 +++++++++++++++++++++ 2 files changed, 26 insertions(+) diff --git a/docs/source/Reduce.md b/docs/source/Reduce.md index 267b8e8..41e43f2 100644 --- a/docs/source/Reduce.md +++ b/docs/source/Reduce.md @@ -121,4 +121,9 @@ Higher values improve search accuracy at the cost of latency. Default: `64` --- +**return_df** (bool, optional) + +When True reduce returns the deduped dataframe rather than clusters + +Default: `False` diff --git a/examples/deduplication.ipynb b/examples/deduplication.ipynb index 5487c79..f62c2d2 100644 --- a/examples/deduplication.ipynb +++ b/examples/deduplication.ipynb @@ -169,6 +169,27 @@ "\n", "print(result.head())" ] + }, + { + "cell_type": "markdown", + "id": "bfa6021a", + "metadata": {}, + "source": [ + "When return_df=True is provided, dt.reduce returns a new Dask DataFrame with semantically similar duplicates removed" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "079d5a86", + "metadata": {}, + "outputs": [], + "source": [ + "deduped_df = dt.reduce(df, action=\"dedup\", embedding_model=\"text-embedding-3-small\", llm=llm, return_df=True)\n", + "\n", + "deduped_df.compute().to_csv(\"deduped_flights.csv\")\n", + "print(deduped_df.head())" + ] } ], "metadata": {