PIPE-RAG: An Execution-Guided Retrieval and Reuse Framework for Reliable Machine Learning Pipeline Code Generation
DOI:
https://doi.org/10.62051/46qs2n55Keywords:
Retrieval-augmented generation; Code reuse; Machine-learning pipelines; Execution feedback; Program repair; Scikit-learn.Abstract
Retrieval-augmented code generation can reduce hallucinated application programming interfaces and encourage reuse, but machine-learning (ML) pipelines impose constraints that ordinary semantic retrieval does not capture: the estimator, task type, data schema, preprocessing order, scoring metric, library version, and executable outputs must agree simultaneously. This paper presents PIPE-RAG, an execution-guided, language-model-compatible framework that treats reusable pipeline code as structured cards rather than untyped text. PIPE-RAG parses a natural-language request into facets, combines word- and character-level retrieval with structured compatibility, applies version-aware maximal marginal relevance, and recomposes a leakage-safe scikit-learn pipeline under explicit constraints. Generated code is parsed, executed in a sandbox, checked for required artifacts, and repaired once when an observed error matches a known API-drift signature. To avoid unverifiable model claims, the reported experiment isolates the deterministic high-confidence retrieval/reuse branch; no proprietary foundation-model API is used. We execute 144 method-task combinations over 36 requests and five real scikit-learn datasets plus one documented mixed-schema stress transformation. PIPE-RAG obtains 100% Recall@1, execution success, and task-complete functional success, compared with 52.78%/41.67% for lexical RAG and 77.78%/58.33% for hybrid RAG. The functional-success gain over hybrid RAG is 41.67 percentage points (paired bootstrap 95% CI: 25.00–58.33; exact McNemar p=6.10×10^-5). A separate seven-case API-drift suite fails before repair and executes successfully after repair in all cases. The results show that structured retrieval, constraint-aware reuse, and execution feedback are complementary controls for dependable ML pipeline synthesis.
Downloads
References
[1] Lewis, P., et al. (2020). Retrieval augmented generation for knowledge intensive NLP tasks. In Advances in Neural Information Processing Systems.
[2] Guu, K., Lee, K., Tung, Z., Pasupat, P., & Chang, M. W. (2020). Retrieval augmented language model pre training. In Proceedings of the ICML (PMLR 119).
[3] Karpukhin, V., et al. (2020). Dense passage retrieval for open domain question answering. In Proceedings of the EMNLP. https://doi.org/10.18653/v1/2020.emnlp-main.550
[4] Borgeaud, S., et al. (2022). Improving language models by retrieving from trillions of tokens. In Proceedings of the ICML (PMLR 162).
[5] Robertson, S., & Zaragoza, H. (2009). The probabilistic relevance framework: BM25 and beyond. Foundations and Trends in Information Retrieval. https://doi.org/10.1561/1500000019
[6] Carbonell, J., & Goldstein, J. (1998). The use of MMR, diversity based reranking for reordering documents and producing summaries. In Proceedings of the SIGIR. https://doi.org/10.1145/290941.291025
[7] Husain, H., Wu, H. H., Gazit, T., Allamanis, M., & Brockschmidt, M. (2019). CodeSearchNet challenge: Evaluating the state of semantic code search. arXiv preprint arXiv:1909.09436.
[8] Feng, Z., et al. (2020). CodeBERT: A pre trained model for programming and natural languages. In Findings of EMNLP. https://doi.org/10.18653/v1/2020.findings-emnlp.139
[9] Guo, D., et al. (2021). GraphCodeBERT: Pre training code representations with data flow. In Proceedings of the ICLR.
[10] Wang, Y., Wang, W., Joty, S., & Hoi, S. C. H. (2021). CodeT5: Identifier aware unified pre trained encoder decoder models for code understanding and generation. In Proceedings of the EMNLP. https://doi.org/10.18653/v1/2021.emnlp-main.685
[11] Ahmad, W. U., Chakraborty, S., Ray, B., & Chang, K. W. (2021). Unified pre training for program understanding and generation. In Proceedings of the NAACL HLT. https://doi.org/10.18653/v1/2021.naacl-main.211
[12] Wang, Y., et al. (2023). CodeT5+: Open code large language models for code understanding and generation. In Proceedings of the EMNLP. https://doi.org/10.18653/v1/2023.emnlp-main.68
[13] Chen, M., et al. (2021). Evaluating large language models trained on code. arXiv preprint arXiv:2107.03374.
[14] Austin, J., et al. (2021). Program synthesis with large language models. arXiv preprint arXiv:2108.07732.
[15] Fried, D., et al. (2023). InCoder: A generative model for code infilling and synthesis. In Proceedings of the ICLR.
[16] Chen, X., Lin, M., Schärli, N., & Zhou, D. (2024). Teaching large language models to self debug. In Proceedings of the ICLR.
[17] Shinn, N., et al. (2023). Reflexion: Language agents with verbal reinforcement learning. In Advances in Neural Information Processing Systems.
[18] Ni, A., et al. (2024). NExT: Teaching large language models to reason about code execution. In Proceedings of the ICML (PMLR 235).
[19] Jain, N., et al. (2021). Jigsaw: Large language models meet program synthesis. arXiv preprint arXiv:2112.02969.
[20] Pedregosa, F., et al. (2011). Scikit learn: Machine learning in Python. Journal of Machine Learning Research.
[21] Thornton, C., Hutter, F., Hoos, H. H., & Leyton Brown, K. (2013). Auto WEKA: Combined selection and hyperparameter optimization of classification algorithms. In Proceedings of the KDD. https://doi.org/10.1145/2487575.2487629
[22] Feurer, M., et al. (2015). Efficient and robust automated machine learning. In Advances in Neural Information Processing Systems.
[23] Olson, R. S., & Moore, J. H. (2016). TPOT: A tree based pipeline optimization tool for automating machine learning. In AutoML Workshop, ICML (PMLR 64).
[24] Olson, R. S., Bartley, N., Urbanowicz, R. J., & Moore, J. H. (2016). Evaluation of a tree based pipeline optimization tool for automating data science. In Proceedings of the GECCO. https://doi.org/10.1145/2908812.2908918
[25] Olson, R. S., & Moore, J. H. (2016). Identifying and harnessing the building blocks of machine learning pipelines for sensible initialization of a data science automation tool. arXiv preprint arXiv:1607.08878.
[26] Kaufman, S., Rosset, S., Perlich, C., & Stitelman, O. (2012). Leakage in data mining: Formulation, detection, and avoidance. ACM Transactions on Knowledge Discovery from Data. https://doi.org/10.1145/2382577.2382579
[27] Kelly, M., Longjohn, R., & Nottingham, K. (n.d.). The UCI machine learning repository. UCI.
[28] Scikit learn Developers. (n.d.). Toy datasets and dataset loading utilities. Scikit learn documentation.
[29] Liu, S., Chen, Y., Xie, X., Siow, J., & Liu, Y. (2020). Retrieval augmented generation for code summarization via hybrid GNN. arXiv preprint arXiv:2006.05405.
[30] Ye, W., et al. (2020). Leveraging code generation to improve code retrieval and summarization via dual learning. In The Web Conference (WWW). https://doi.org/10.1145/3366423.3380295
[31] Zhao, W., Chen, T., Yang, J. S., & Qiu, L. (2026). AutoML Pipeline: A RAG enhanced code generation framework with pre validation for cloud native machine learning workflows. IEEE Access, 14, 41932–41945. https://doi.org/10.1109/ACCESS.2026.3673923
[32] Mo, T., Zhang, C., Zou, J., Guo, Z., & Rhee, M. (2026). Self evolving AI agents with dual memory for automated software testing and bug localization. IEEE Access, 14, 111086–111102. https://doi.org/10.1109/ACCESS.2026.3713401
[33] Rhee, M., Zou, J., Mo, T., Teng, D., & Yang, J. S. (2026). Enhancing web search agents with self play contrastive fine tuning. IEEE Access. https://doi.org/10.1109/ACCESS.2026.3717594
[34] Teng, D., Rhee, M., Qin, Y., Zi, B., & Liu, W. (2026). SW SpeedDLM: Sliding window speculative decoding for diffusion language models under long context constraints. Mathematics, 14(12), 2137. https://doi.org/10.3390/math14122137
Downloads
Published
Issue
Section
License
Copyright (c) 2026 Transactions on Computer Science and Intelligent Systems Research
This work is licensed under a Creative Commons Attribution-NonCommercial 4.0 International License.
