Grounding Data Science Code Generation with Input-Output Specifications

TL;DR

Gift4Code tackles misalignment between natural language prompts and data-science code outputs by grounding NL intents with execution-derived I/O specifications. The method generates synthetic NL intents and code in realistic programmatic contexts, executes candidate solutions to infer I/O constraints, and then fine-tunes a code LLM to produce solutions conditioned on intents augmented with I/O specifications. Empirical results on Arcade and DS-1000 show substantial gains in both executability and alignment to user specifications, with I/O summaries delivering the strongest improvements. The approach demonstrates how execution feedback can be leveraged to produce more reliable data-science code generation and is extensible to other domains demanding precise task specifications.

Abstract

Large language models (LLMs) have recently demonstrated a remarkable ability to generate code from natural language (NL) prompts. However, in the real world, NL is often too ambiguous to capture the true intent behind programming problems, requiring additional input-output (I/O) specifications. Unfortunately, LLMs can have difficulty aligning their outputs with both the NL prompt and the I/O specification. In this paper, we give a way to mitigate this issue in the context of data science programming, where tasks require explicit I/O specifications for clarity. Specifically, we propose GIFT4Code, a novel approach for the instruction fine-tuning of LLMs with respect to I/O specifications. Our method leverages synthetic data produced by the LLM itself and utilizes execution-derived feedback as a key learning signal. This feedback, in the form of program I/O specifications, is provided to the LLM to facilitate instruction fine-tuning. We evaluated our approach on two challenging data science benchmarks, Arcade and DS-1000. The results demonstrate a significant improvement in the LLM's ability to generate code that is not only executable but also accurately aligned with user specifications, substantially improving the quality of code generation for complex data science tasks.

Figures (9)

• Figure 1: Left: Illustration of how users prompt code LLMs with NL intents and I/O specifications to generate code with complex outputs (pandasDataframes). Vanilla code LLMs fail to understand I/O specifications. Right: Our proposed instruction tuning approach uses synthetic intents and code solutions, where intents are augmented with I/O specifications derived from program execution. Models trained on such synthetic data could better follow a user's intent.
• Figure 2: Two examples on Arcade. Left: the base code LLM does not follow the intent to group the "League" column. Its few-shot prompting variant failed to correctly utilize the specifications. Gift4Code's solution aligns to the user's intent. Right: code LLM tries to call a non-exsiting column, leading to a KeyError. Its few-shot variant follows the specifications incorrectly, leading to the syntax error. Gift4Code generates correct solution. The prompt always includes dataframe schema, regardless of whether specifications are provided.
• Figure 3: Frequency of error types and code execution rate for different methods. Bottom x-axis stands for the counts of schema errors and syntax errors. Top x-axis represents execution rate. Instruction fine-tuning without specifications (SyntheticFT) decreases executability. Few-shot prompting with specifications reduces schema understanding errors with more syntax errors. Gift4Code achieves the best performance by combining their strengths.
• Figure : An exemplar in the few-shot prompt when querying LLM to generate NL intents.
• Figure : Programmatic context and generated questions on an example concerning pokemon game experience dataset.