Implementation:FMInference FlexLLMGen Batch Query Test

Field	Value
Sources	FlexLLMGen
Domains	Batch_Processing, Data_Wrangling
Last Updated	2026-02-09 00:00 GMT

Overview

Concrete tool for running batched LLM inference over data wrangling datasets provided by the FlexLLMGen data wrangling application.

Description

batch_query_test() takes parsed args, iterates over test examples in batches, constructs prompts by combining few-shot prefix with serialized query text, tokenizes and pads to uniform length, creates Policy and OptLM per batch (since sequence length may vary), runs OptLM.generate(), decodes output, extracts predicted labels by stripping the input prefix, and collects predictions for metric evaluation. Results are written to .feather files.

Usage

Called internally by data_wrangle_run.py when --batch_run is set. Invoked via CLI.

Code Reference

Source: flexllmgen/apps/data_wrangle/data_wrangle_run.py, Lines: 324-489
Import: Internal function, invoked via CLI: python -m flexllmgen.apps.data_wrangle.data_wrangle_run --data_dir ... --model ... --batch_run

Signature:

def batch_query_test(
    args,           # Parsed CLI arguments
    data,           # Dict of DataFrames from read_data()
    prompt_prefix,  # Few-shot prompt string
    log_dir,        # Output directory for results
):
    """Run batched inference over test dataset.

    Args:
        args: CLI args with model, percent, path, etc.
        data: Dict with "train", "test" DataFrames
        prompt_prefix: Few-shot examples string
        log_dir: Directory for saving results
    """

I/O Contract

Direction	Name	Type	Required	Description
Input	args	argparse.Namespace	Yes	CLI arguments
Input	data	Dict[str, DataFrame]	Yes	Loaded dataset
Input	prompt_prefix	str	Yes	Few-shot prompt
Input	log_dir	str	Yes	Output directory
Output	preds	List[str]	--	Predictions extracted from generation
Output	results	.feather / .json	--	Saved to {log_dir}/trial_{n}.feather and metrics.json

Usage Examples

# CLI invocation (primary interface)
python -m flexllmgen.apps.data_wrangle.data_wrangle_run \
    --data_dir fm_data_tasks/data/entity_matching/structured/Amazon-Google \
    --model facebook/opt-30b \
    --path ~/opt_weights \
    --percent 0 100 0 100 100 0 \
    --batch_run \
    --k 5 \
    --sample_method manual \
    --num_run 100

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment