Map Operation

The Map operation applies a specified transformation in natural language to each item in a column of input data.

Core Implementation

nirvana.ops.map.MapOperation(user_instruction: str = '', input_columns: list[str] = [], output_columns: list[str] = [], context: list[dict] | str | None = None, model: str | None = None, tool: Callable | BaseTool | None = None, strategy: Literal['plain', 'fewshot', 'self-refine'] = 'plain', limit: int | None = None, rate_limit: int = 16, assertions: list[Callable] | None = [])

Bases: BaseOperation

Map operator: Applies an LLM to perform a transformation on a column, producing a new column as output

Source code in nirvana/ops/map.py

def __init__(
    self,
    user_instruction: str = "",
    input_columns: list[str] = [],
    output_columns: list[str] = [],
    context: list[dict] | str | None = None,
    model: str | None = None,
    tool: Callable | BaseTool | None = None,
    strategy: Literal["plain", "fewshot", "self-refine"] = "plain",
    limit: int | None = None,
    rate_limit: int = 16,
    assertions: list[Callable] | None = [],
):
    if tool and not isinstance(tool, BaseTool):
        tool = FunctionCallTool.from_function(func=tool)

    super().__init__(
        op_name="map",
        user_instruction=user_instruction,
        context=context,
        model=model,
        tool=tool,
        strategy=strategy,
        limit=limit,
        rate_limit=rate_limit,
        assertions=assertions,
    )
    self.prompter = MapPrompter()
    self.input_columns = input_columns
    self.output_columns = output_columns

Attributes

strategy_options = ['plain', 'fewshot', 'self-refine'] class-attribute instance-attribute

prompter = MapPrompter() instance-attribute

input_columns = input_columns instance-attribute

output_columns = output_columns instance-attribute

dependencies: list[str] property

generated_fields: list[str] property

op_kwargs property

Functions

execute(input_data: pd.DataFrame, **kwargs) async

Source code in nirvana/ops/map.py

async def execute(
    self, 
    input_data: pd.DataFrame,
    **kwargs
):  
    if self.user_instruction is None and not self.has_udf():
        raise ValueError("Neither `user_instruction` nor `func` is given.")

    if input_data.empty:
        return MapOpOutputs(outputs={}, cost=0.0)

    processed_data = input_data[self.input_columns]
    dtypes = []
    for col in self.input_columns:
        if isinstance(input_data[col].dtype, ImageDtype):
            dtypes.append("image")
        else:
            dtypes.append("text")

    if self.strategy == "plain":
        execution_func = functools.partial(self._execute_by_plain_llm, dtypes=dtypes, model=self.model, **kwargs)
    elif self.strategy == "fewshot":
        assert self.context is not None, "Few-shot examples must be provided in the context for in-context learning."
        demos = self.context
        execution_func = functools.partial(self._execute_by_fewshot_llm, dtypes=dtypes, demos=demos, model=self.model, **kwargs)
    elif self.strategy == "self_refine":
        execution_func = functools.partial(self._execute_by_self_refine, dtypes=dtypes, model=self.model, **kwargs)
    else:
        raise ValueError(f"The optional strategies available for map are {self.strategy_options}. Strategy {self.strategy} is not supported.")

    # Create tasks for all data points
    tasks = []
    empty_map_result = {column: None for column in self.output_columns}
    for _, data in processed_data.iterrows():
        if data.empty:
            tasks.append(asyncio.create_task(asyncio.sleep(0, result=(empty_map_result, 0.0))))
        elif self.has_udf():
            tasks.append(asyncio.create_task(self._execute_by_func(data, self.user_instruction, self.tool, execution_func)))
        else:
            tasks.append(asyncio.create_task(execution_func(data, self.user_instruction)))

    # Wait for all tasks to complete
    if self.limit is not None and self.limit <= 0:
        warnings.warn("The limit should be positive. To execute, the limit will be ignored.")
        self.limit = None

    token_cost: float = 0.0
    map_outputs: dict[str, list] = {column: [] for column in self.output_columns}
    if self.limit is not None:
        num_passed_records: int = 0
        reach_limit: bool = False
        for i in range(0, len(tasks), self.limit):
            if reach_limit:
                break
            batch_tasks = tasks[i:i + self.limit]
            batch_results = await asyncio.gather(*batch_tasks)
            token_cost += sum([result[1] for result in batch_results])
            for result, _ in batch_results:
                num_passed_records += 1
                for column in self.output_columns:
                    map_outputs[column].append(result[column])
                if num_passed_records >= self.limit:
                    reach_limit = True
                    break
        num_remaining_records = len(processed_data) - num_passed_records
        if num_remaining_records > 0:
            for column in self.output_columns:
                map_outputs[column].extend([None] * num_remaining_records)
    else:
        results = await asyncio.gather(*tasks)
        token_cost += sum([result[1] for result in results])
        for result, _ in results:
            for column in self.output_columns:
                map_outputs[column].append(result[column])

    return MapOpOutputs(
        outputs=map_outputs,
        cost=token_cost
    )

Output Class

nirvana.ops.map.MapOpOutputs(cost: float = 0.0, outputs: dict[str, list] = dict()) dataclass

Bases: BaseOpOutputs

Attributes

outputs: dict[str, list] = field(default_factory=dict) class-attribute instance-attribute

Functions

__add__(other: MapOpOutputs)

Source code in nirvana/ops/map.py

def __add__(self, other: "MapOpOutputs"):
    assert self.outputs.keys() == other.outputs.keys(), "Cannot merge MapOpOutputs with different field names."
    new_output = {}
    for field_name in self.outputs.keys():
        new_output[field_name] = self.outputs[field_name] + other.outputs[field_name]
    return MapOpOutputs(
        outputs=new_output,
        cost=self.cost + other.cost
    )

Function Wrapper

nirvana.ops.map

map_wrapper(input_data: DataFrame, user_instruction: str = None, input_columns: list[str] = None, output_columns: list[str] = None, context: list[dict] | str | None = None, model: str | None = None, func: Callable = None, strategy: Literal['plain', 'fewshot', 'self-refine'] = 'plain', limit: int | None = None, rate_limit: int = 16, assertions: list[Callable] | None = [], **kwargs)

A function wrapper for map operation

Parameters:

Name	Type	Description	Default
input_data	DataFrame	Input dataframe	required
user_instruction	str	User instruction. Defaults to None.	None
input_columns	list[str]	Input column. Defaults to None.	None
output_columns	list[str]	Output columns. Defaults to None.	None
context	list[dict] | str	Context. Defaults to None.	None
model	str	Model. Defaults to None.	None
func	Callable	User function. Defaults to None.	None
strategy	Literal['plain', 'fewshot', 'self-refine']	Strategy. Defaults to "plain".	'plain'
limit	int	Maximum number of outputs to produce before stopping.	None
rate_limit	int	Rate limit. Defaults to 16.	16
assertions	list[Callable]	Assertions. Defaults to [].	[]
**kwargs		Additional keyword arguments for OpenAI Clent.	{}

Source code in nirvana/ops/map.py

def map_wrapper(
    input_data: pd.DataFrame, 
    user_instruction: str = None,
    input_columns: list[str] = None,
    output_columns: list[str] = None,
    context: list[dict] | str | None = None,
    model: str | None = None,
    func: Callable = None,
    strategy: Literal["plain", "fewshot", "self-refine"] = "plain",
    limit: int | None = None,
    rate_limit: int = 16,
    assertions: list[Callable] | None = [],
    **kwargs
):
    """
    A function wrapper for map operation

    Args:
        input_data (pd.DataFrame): Input dataframe
        user_instruction (str, optional): User instruction. Defaults to None.
        input_columns (list[str], optional): Input column. Defaults to None.
        output_columns (list[str], optional): Output columns. Defaults to None.
        context (list[dict] | str, optional): Context. Defaults to None.
        model (str, optional): Model. Defaults to None.
        func (Callable, optional): User function. Defaults to None.
        strategy (Literal["plain", "fewshot", "self-refine"], optional): Strategy. Defaults to "plain".
        limit (int): Maximum number of outputs to produce before stopping.
        rate_limit (int, optional): Rate limit. Defaults to 16.
        assertions (list[Callable], optional): Assertions. Defaults to [].
        **kwargs: Additional keyword arguments for OpenAI Clent.
    """

    map_op = MapOperation(
        user_instruction=user_instruction,
        strategy=strategy,
        input_columns=input_columns,
        output_columns=output_columns,
        context=context,
        model=model,
        tool=FunctionCallTool.from_function(func=func) if func else None,
        limit=limit,
        rate_limit=rate_limit,
        assertions=assertions
    )
    outputs = asyncio.run(map_op.execute(
        input_data=input_data,
        **kwargs
    ))
    return outputs