D oc CG en: Document-based Controlled Code Generation (2024)

Abstract: Recent developments show that Large Language Models (LLMs) produce state-of-the-art performance on natural language (NL) to code generation for resource-rich general-purpose languages like C++, Java, and Python. However, their practical usage for structured domain-specific languages (DSLs) such as YAML, JSON is limited due to domain-specific schema, grammar, and customizations generally unseen by LLMs during pre-training. Efforts have been made to mitigate this challenge via in-context learning through relevant examples or by fine-tuning. However, it suffers from problems, such as limited DSL samples and prompt sensitivity but enterprises maintain good documentation of the DSLs. Therefore, we propose DocCGen, a framework that can leverage such rich knowledge by breaking the NL-to-Code generation task for structured code languages into a two-step process. First, it detects the correct libraries using the library documentation that best matches the NL query. Then, it utilizes schema rules extracted from the documentation of these libraries to constrain the decoding. We evaluate our framework for two complex structured languages, Ansible YAML and Bash command, consisting of two settings: Out-of-domain (OOD) and In domain (ID). Our extensive experiments show that DocCGen consistently improves different sized language models across all six evaluation metrics, reducing syntactic and semantic errors in structured code.

## introduction the natural language to code (nl-to-code) task has become pivotal in the intersection of natural language processing and programming. nl-to-code systems can help engineers write a program efficiently by conveying their intentions at a higher level, as shown in figure 1 . systems like amazon code whisperer 2 , github co-pilot 3 perform well in nl-to-code task due to large language models (llm) trained on extensive data. while they perform well in general resource-rich languages like c++, python, or java, their practical usage in structured dsl is limited. dsls are enterprise-specific languages with specialized schemas and syntax suitable for a specific domain or application 4 . numerous enterprises use structured languages like bash, yaml, json and hcl (hashicorp configuration language) with specific customizations for automation and to configure and manage infrastructure in it environments. these languages or their customizations are potentially unseen by lms during pre-training, limiting their practical usage @xcite . some existing methods attempt to address this challenge via in-context learning through examples @xcite , by fine-tuning @xcite or by using relevant documentation as additional context @xcite @xcite . however, relevant context or samples available for dsl are often insufficient to incorporate diverse library schema rules or specialized structure knowledge in the lm @xcite . this results in hallucination and different syntactic and semantic errors, as shown in figure 1 . however, enterprises usually maintain detailed documentation of their custom libraries (e.g. ansible modules, bash utilities), including the descriptions, schema, and syntax, to assist developers in enforcing structure and maintaining data integrity. we believe such schema and documentation can be better leveraged during code generation. therefore, we propose a framework doccgen that treats the nl-to-code task as a two-step process, each heavily relying on the documentation. the first step identifies relevant code libraries for the task by retrieving the library documentation relevant to the nl query. the second step employs constrained @xcite to guide code generation by using the grammar and schema rules extracted from the documentation of libraries identified in the first step, as shown in figure 2 . we evaluate this approach for two diverse and complex structured languages, ansible yaml and bash command. generation for these languages is tricky due to complexities like the diverse library schemas, optional and required fields, the order-agnostic nature of fields, and inter-field dependencies. we believe studying these complex structures encompasses most of the challenges in other structured dsls and allows easily extending doccgen to other domains. since the major challenge in dsls is the limited availability of samples, we focus on enhancing performance for unseen code libraries or libraries with very few samples in the training corpus. hence, we evaluate our approach in two settings: in-domain and outof-domain. similar to @xcite , none of the libraries in the test set are seen during training in the ood setting. in the id setting, every library in the test set has very few nl-to-code pairs in the train set. doccgen consistently improve over stateof-the-art models and techniques by a significant margin ( finally, we introduce first publicly available benchmark dataset for nl to structured code generation task consisting of ansible-yaml language. intricate challenges in ansible-yaml generation, like the complex structure and diverse module schemas, lead to subpar performance even for finetuned code lms ( ## doccgen framework doccgen is a two-stage framework: the first stage uses information retrieval (ir) to detect relevant libraries. the second stage uses the neuro-symbolic constrained decoding to control generation and ensure adherence to the schema of relevant libraries. ## dataset this section describes datasets for nl to bash and ansible yaml task, including augmenting datasets with module descriptions and schema information. ## experiments in this section, we lay out our experiments across nl-to-code tasks and datasets. ## results and analysis results and comparison of our framework with various baselines are presented in tables 1, 2 and 3. this section presents several observations and a qualitative analysis of the performance. improvement in module accuracy: we observe that extended pre-training does not improve performance in structured dsls with limited code samples in the documentation. therefore, we use an ir-based approach that focuses on retrieving utility descriptions, unlike @xcite , which retrieves passages with options (flags and sub-commands) and utilities. this targeted detection reduces the search space for ir from 400k to 1.5k documents, leading to a notable improvement in hits@1 ( the model to follow one of the templates for k retrieved libraries. cmd acc and module acc drop with a higher value of k ( ## conclusion we propose doccgen, a novel framework for nlto-code generation for structured dsls. doccgen decomposes the nl-to-code generation into two steps involving the detection of relevant libraries in the first step and using schema and grammar rules extracted from the documentation of these libraries to guide the decoding in the second step. we evaluate the performance of doccgen for two complex structured languages, bash command and ansible yaml, involving two settings, ood and id. our approach outperforms state-of-the-art techniques consistently across all metrics for differentsized models. it reduces syntactic and semantic errors in code, particularly for unseen libraries and low-resource data settings. we also contribute the first publicly available benchmark dataset for nl to ansible-yaml task. we augment nl to ansible-yaml and tldr dataset with description and schema information. we hope this work will help advance research in solving dsl-related tasks and constrained generation.