Inference API Reference

LLMSQL Transformers Inference Function

This module provides a single function inference_transformers() that performs text-to-SQL generation using large language models via the Transformers backend.

Example

from llmsql.inference import inference_transformers

results = inference_transformers(
    model_or_model_name_or_path="Qwen/Qwen2.5-1.5B-Instruct",
    repo_id="llmsql-bench/llmsql-2.0",
    output_file="outputs/preds_transformers.jsonl",
    num_fewshots=5,
    batch_size=8,
    max_new_tokens=256,
    temperature=0.7,
    model_kwargs={
        "torch_dtype": "bfloat16",
    },
    generation_kwargs={
        "do_sample": False,
    },
)

Notes

This function uses the HuggingFace Transformers backend and may produce slightly different outputs than the vLLM backend even with the same inputs due to differences in implementation and numerical precision.

llmsql.inference.inference_transformers.inference_transformers(model_or_model_name_or_path: str | AutoModelForCausalLM, tokenizer_or_name: str | Any | None = None, *, trust_remote_code: bool = True, dtype: dtype | str = torch.float16, device_map: str | dict[str, int] | None = 'auto', hf_token: str | None = None, model_kwargs: dict[str, Any] | None = None, tokenizer_kwargs: dict[str, Any] | None = None, chat_template: str | None = None, max_new_tokens: int = 256, temperature: float = 0.0, do_sample: bool = False, top_p: float = 1.0, top_k: int = 50, generation_kwargs: dict[str, Any] | None = None, version: Literal['1.0', '2.0'] = '2.0', output_file: str = 'llm_sql_predictions.jsonl', workdir_path: str | None = None, num_fewshots: int = 5, batch_size: int = 8, limit: int | float | None = None, seed: int = 42) → list[dict[str, str]][source]

Inference a causal model (Transformers) on the LLMSQL benchmark.

Parameters:

  • model_or_model_name_or_path – Model object or HF model name/path.

  • tokenizer_or_name – Tokenizer object or HF tokenizer name/path.

  • Loading (# Tokenizer)

  • trust_remote_code – Whether to trust remote code (default: True).

  • dtype – Torch dtype for model (default: float16). Strings such as “bfloat16” or “auto” are accepted as well.

  • device_map – Device placement strategy (default: “auto”).

  • hf_token – Hugging Face authentication token.

  • model_kwargs – Additional arguments for AutoModelForCausalLM.from_pretrained(). Values here take precedence over ‘dtype’, ‘device_map’, ‘trust_remote_code’ and ‘token’. A ‘dtype’ (or legacy ‘torch_dtype’) entry replaces the dtype argument and may be given as a string.

  • Loading

  • tokenizer_kwargs – Additional arguments for AutoTokenizer.from_pretrained(). ‘padding_side’ defaults to “left”. Note: ‘trust_remote_code’, ‘token’ are handled separately and will override values here.

  • Chat (# Prompt &)

  • chat_template – Optional chat template to apply before tokenization.

  • Generation (#)

  • max_new_tokens – Maximum tokens to generate per sequence.

  • temperature – Sampling temperature (0.0 = greedy).

  • do_sample – Whether to use sampling vs greedy decoding.

  • top_p – Nucleus sampling parameter.

  • top_k – Top-k sampling parameter.

  • generation_kwargs – Additional arguments for model.generate(). Note: ‘max_new_tokens’, ‘temperature’, ‘do_sample’, ‘top_p’, ‘top_k’ are handled separately.

  • Benchmark (#)

  • version – LLMSQL version

  • output_file – Output JSONL file path for completions.

  • workdir_path – Directory to store downloaded benchmark files. If omitted, a temporary directory is created automatically.

  • num_fewshots – Number of few-shot examples (0, 1, or 5).

  • batch_size – Batch size for inference.

  • seed – Random seed for reproducibility.

  • limit – Limit the number of questions to evaluate. If an integer, evaluates the first N samples. If a float between 0.0 and 1.0, evaluates the first X*100% of samples. If None, evaluates all samples (default).

Returns:

List of generated SQL results with metadata.

—

LLMSQL vLLM Inference Function

This module provides a single function inference_vllm() that performs text-to-SQL generation using large language models via the vLLM backend.

Example

from llmsql.inference import inference_vllm

results = inference_vllm(
    model_name="Qwen/Qwen2.5-1.5B-Instruct",
    version="2.0",
    tables_path="data/tables.jsonl",
    num_fewshots=5,
    batch_size=8,
    max_new_tokens=256,
    temperature=0.7,
    tensor_parallel_size=1,
    lora_path="path/to/lora"
)

Notes

This function uses the vLLM backend. Outputs may differ from the Transformers backend due to differences in implementation, batching, and numerical precision.

llmsql.inference.inference_vllm.inference_vllm(model_name: str, *, trust_remote_code: bool = True, tensor_parallel_size: int = 1, hf_token: str | None = None, llm_kwargs: dict[str, Any] | None = None, use_chat_template: bool = True, lora_config: dict[str, Any] | None = None, max_new_tokens: int = 256, temperature: float = 1.0, do_sample: bool = True, sampling_kwargs: dict[str, Any] | None = None, version: Literal['1.0', '2.0'] = '2.0', output_file: str = 'llm_sql_predictions.jsonl', workdir_path: str | None = None, limit: int | float | None = None, num_fewshots: int = 5, batch_size: int = 8, seed: int = 42) → list[dict[str, str]][source]

Run SQL generation using vLLM.

Parameters:

  • model_name – Hugging Face model name or path.

  • Loading (# Model)

  • trust_remote_code – Whether to trust remote code (default: True).

  • tensor_parallel_size – Number of GPUs for tensor parallelism (default: 1).

  • hf_token – Hugging Face authentication token.

  • llm_kwargs – Additional arguments for vllm.LLM(). Note: ‘model’, ‘tokenizer’, ‘tensor_parallel_size’, ‘trust_remote_code’ are handled separately and will override values here.

  • lora_config –

    Optional dict with LoRA parameters: - lora_path: Path to the pretrained LoRA adapter (required if enable_lora) - lora_name: Logical name for the LoRA adapter - lora_scale: Scaling factor for LoRA weights - max_lora_rank: Maximum LoRA rank supported by vLLM

    LoRA usage rules:

    • If lora_config is provided, enable_lora must be True in llm_kwargs.

    • If enable_lora is True, a valid lora_config must be provided.

    • Otherwise, an exception is raised to prevent inconsistent configuration.

  • Generation (#)

  • max_new_tokens – Maximum tokens to generate per sequence.

  • temperature – Sampling temperature (0.0 = greedy).

  • do_sample – Whether to use sampling vs greedy decoding.

  • sampling_kwargs – Additional arguments for vllm.SamplingParams(). Note: ‘temperature’, ‘max_tokens’ are handled separately and will override values here.

  • Benchmark (#)

  • version – LLMSQL version

  • output_file – Path to write outputs (will be overwritten).

  • workdir_path – Directory to store downloaded benchmark files. If omitted, a temporary directory is created automatically.

  • num_fewshots – Number of few-shot examples (0, 1, or 5).

  • batch_size – Number of questions per generation batch.

  • seed – Random seed for reproducibility.

  • limit – Limit the number of questions to evaluate. If an integer, evaluates the first N samples. If a float between 0.0 and 1.0, evaluates the first X*100% of samples. If None, evaluates all samples (default).

Returns:

List of dicts containing question_id and generated completion.

—

LLMSQL OpenAI-Compatible API Inference Function

This module provides inference_api() for text-to-SQL generation against an OpenAI-compatible Chat Completions API.

llmsql.inference.inference_api.inference_api(model_name: str, *, base_url: str, endpoint: str = 'chat/completions', api_key: str | None = None, timeout: float = 120.0, requests_per_minute: float | None = None, api_kwargs: dict[str, Any] | None = None, request_headers: dict[str, str] | None = None, version: Literal['1.0', '2.0'] = '2.0', output_file: str = 'llm_sql_predictions.jsonl', workdir_path: str | None = None, limit: int | float | None = None, num_fewshots: int = 5, seed: int = 42) → list[dict[str, str]][source]

Run SQL generation using an OpenAI-compatible Chat Completions API.

Requests are dispatched concurrently so that HTTP round-trip time does not count against the rate-limit interval — achieving a true requests_per_minute throughput rather than requests_per_minute / (1 + latency_in_minutes).

Parameters:

  • model_name – The model name of the api.

  • base_url – e.g. “https://api.openai.com/v1/”

  • endpoint – e.g. “chat/completions”

  • Benchmark (#)

  • version – LLMSQL version

  • output_file – Path to write outputs (will be overwritten).

  • workdir_path – Directory to store downloaded benchmark files. If omitted, a temporary directory is created automatically.

  • num_fewshots – Number of few-shot examples (0, 1, or 5).

  • batch_size – Number of questions per generation batch.

  • seed – Random seed for reproducibility.

  • limit – Limit the number of questions to evaluate. If an integer, evaluates the first N samples. If a float between 0.0 and 1.0, evaluates the first X*100% of samples. If None, evaluates all samples (default).

Returns:

List of dicts containing question_id and generated completion.

—

LLMSQL Custom Function Inference

This module provides inference_function() for text-to-SQL generation using an arbitrary user-provided async inference callable.

llmsql.inference.inference_function.inference_function(*, inference_function: Callable[[...], Awaitable[str]], requests_per_minute: float | None = None, max_concurrency: int | None = 32, raise_on_error: bool = False, function_kwargs: dict[str, Any] | None = None, version: Literal['1.0', '2.0'] = '2.0', output_file: str = 'llm_sql_predictions.jsonl', workdir_path: str | None = None, limit: int | float | None = None, num_fewshots: int = 5, seed: int = 42) → list[dict[str, str]][source]

Run SQL generation using a user-provided async callable.

LLMSQL downloads the benchmark, builds the prompt for every question (with the requested number of few-shot examples) and awaits your callable for each of them. This lets you plug in any engine, API client, router or agent while keeping the standard LLMSQL prompts and output format, so the resulting file can be passed directly to llmsql.evaluate().

The callable is awaited as:

await inference_function(
    prompt,                      # str, the fully built LLMSQL prompt
    question=question,           # dict, the raw benchmark question row
    table=table,                 # dict, the table the question refers to
    **function_kwargs,
)

and must return the model completion (it is converted with str()). Calls run concurrently on a single event loop, bounded by max_concurrency and spaced out by requests_per_minute. Results are appended to output_file as soon as each call finishes, so the file order follows completion order, not question order.

Error handling: by default (raise_on_error=False) an exception raised by the callable is logged together with the question_id and an empty completion is recorded for that question (it will be counted as incorrect by the evaluator); a summary of failures is logged at the end. With raise_on_error=True the first exception cancels all remaining calls and is re-raised. A callable that does not return an awaitable always raises TypeError.

The function can be called from synchronous code as well as from within a running event loop (e.g. Jupyter), in which case nest_asyncio is applied to that loop.

Example

>>> from llmsql import inference_function
>>> async def my_model(prompt, **kwargs):
...     return "SELECT 1"
>>> results = inference_function(
...     inference_function=my_model,
...     requests_per_minute=60,
...     max_concurrency=8,
... )
Parameters:

  • inference_function – Async callable (async def) that receives the prompt as the first positional argument plus question, table and **function_kwargs keyword arguments, and returns the generated SQL completion.

  • requests_per_minute – Maximum number of calls started per minute. If None (default), calls are not rate limited.

  • max_concurrency – Maximum number of calls in flight at the same time. Defaults to 32. Use None to disable the cap (not recommended for the full benchmark, as all questions would be dispatched at once).

  • raise_on_error – If True, re-raise the first exception raised by the callable and abort the run. If False (default), log the error and record an empty completion for that question.

  • function_kwargs – Extra keyword arguments forwarded to every call, e.g. sampling parameters such as {"temperature": 0.0}.

  • version – LLMSQL benchmark version ("1.0" or "2.0").

  • output_file – Path of the JSONL file to write outputs to (it is overwritten).

  • workdir_path – Directory to store downloaded benchmark files. If omitted, a temporary directory is created automatically.

  • limit – Limit the number of questions to evaluate. If an integer, evaluates the first N samples. If a float between 0.0 and 1.0, evaluates the first X*100% of samples. If None, evaluates all samples (default).

  • num_fewshots – Number of few-shot examples (0, 1, or 5).

  • seed – Random seed for reproducibility.

Returns:

List of dicts containing question_id and generated completion, in completion order.

Raises:

  • TypeError – If inference_function is not callable or does not return an awaitable.

  • ValueError – If requests_per_minute, max_concurrency or limit has an invalid value.

  • Exception – Any exception raised by inference_function when raise_on_error=True.

—

💬 Made with ❤️ by the LLMSQL Team