diff --git a/datatune/__init__.py b/datatune/__init__.py index e66a30e..0f288b8 100644 --- a/datatune/__init__.py +++ b/datatune/__init__.py @@ -2,4 +2,6 @@ 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 SemanticDeduplicator +from datatune.core.reduce import reduce +__all__ = ["map", "filter", "finalize", "Agent", "reduce"] \ No newline at end of file diff --git a/datatune/core/dask/filter_dask.py b/datatune/core/dask/filter_dask.py index 1d39629..56b0255 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, + clusters: 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 clusters + 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, + clusters: 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.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__" @@ -240,6 +266,7 @@ def __call__(self, llm: Callable, df: Dict): self.llm_output_column, self.prompt, self.serialized_input_column, + self.clusters, ), meta=meta_dict, ) diff --git a/datatune/core/dask/map_dask.py b/datatune/core/dask/map_dask.py index 0d5dfa1..deb0a43 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() ] @@ -40,6 +41,7 @@ def llm_batch_inference( prompt: str, serialized_input_column: str, expected_new_fields: List[str], + clusters: List[dict], df: pd.DataFrame, ) -> pd.DataFrame: """ @@ -71,15 +73,36 @@ 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] - df[llm_output_column] = llm( - df[serialized_input_column], prefix, prompt, suffix, optimized=True + dup_to_canon = { + dup: c["canonical_id"] + for c in clusters + 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 +213,13 @@ def __init__( input_fields: Optional[List] = None, output_fields: Optional[List] = None, name: Optional[str] = None, + clusters: List[dict] = None, ): super().__init__(name=name) self.prompt = prompt self.input_fields = input_fields self.output_fields = output_fields + 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__" @@ -234,6 +259,7 @@ def __call__(self, llm: Callable, df: Dict): self.prompt, self.serialized_input_column, self.output_fields, + self.clusters ), meta=meta_dict, ) diff --git a/datatune/core/deduplication.py b/datatune/core/deduplication.py new file mode 100644 index 0000000..f123a1f --- /dev/null +++ b/datatune/core/deduplication.py @@ -0,0 +1,389 @@ +import glob +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 +import dask +import os +from .reduce import register_action + +logger = get_logger(__name__) +def input_as_string( + serialized_input_column: str, + df: pd.DataFrame, +) -> 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 = df.astype(object).where(df.notna(), None) + + df[serialized_input_column] = [ + str(row.to_dict()) for _, row in df.iterrows() + ] + return df + +@register_action("dedup") +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, + return_df: bool = False, + +): + 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 + self.return_df = return_df + + + def _embed_and_write_partition( + self, + part, # lazy Dask Series + partition_id: int, + output_dir: str, + ): + pdf = part + + if pdf.empty: + print("empty partition") + return 0 + + 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): + 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 + + 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 _llm_evaluation(self, clusters, df_column): + """ + 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.""" + "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 a string "NO_DUPLICATES" enclosed in double quotes. + + Do not add explanations. + Do not reference other clusters. + """ + from collections import defaultdict + all_needed_indices = set(idx for cluster in clusters for idx in cluster) + cluster_rows = defaultdict(list) + for pid in range(df_column.npartitions): + part = df_column.get_partition(pid).compute() + 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 + 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: + if idx in row_map: + lines.append(f"ID {idx}: {row_map[idx]}") + lines.append(f"CLUSTER {cluster_id} END") + llm_inputs.append("\n".join(lines)) + + logger.info("Sending %d clusters to LLM for evaluation...", len(llm_inputs)) + llm_outputs = self.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 + elif isinstance(output, set): + 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 _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) + + df = df.map_partitions( + partial( + input_as_string, + "serialized_input_column", + ) + ) + #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, + ) + llm_outputs = self._llm_evaluation( + clusters, df["serialized_input_column"] + ) + + 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 + + diff --git a/datatune/core/filter.py b/datatune/core/filter.py index 2421fd1..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): +def filter(*, prompt, input_fields=None, clusters=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, + 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 a5113ef..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): +def map(*, prompt, output_fields, input_fields=None, clusters=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, + clusters=clusters )(llm, data) elif _is_ibis_table(data): from .ibis.map_ibis import _map_ibis diff --git a/datatune/core/reduce.py b/datatune/core/reduce.py new file mode 100644 index 0000000..3879279 --- /dev/null +++ b/datatune/core/reduce.py @@ -0,0 +1,22 @@ +import dask.dataframe as dd + + +_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 diff --git a/docs/source/Reduce.md b/docs/source/Reduce.md new file mode 100644 index 0000000..41e43f2 --- /dev/null +++ b/docs/source/Reduce.md @@ -0,0 +1,129 @@ +# 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` + +--- +**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 new file mode 100644 index 0000000..f62c2d2 --- /dev/null +++ b/examples/deduplication.ipynb @@ -0,0 +1,202 @@ +{ + "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())" + ] + }, + { + "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": { + "language_info": { + "name": "python" + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/pyproject.toml b/pyproject.toml index 3847728..f8d1ee2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -17,7 +17,8 @@ dependencies = [ "numpy", "litellm", "dask[dataframe]", - "pandas" + "pandas", + "faiss-cpu" ] [project.optional-dependencies]