← Back to Volume 13

PIPE-RAG: An Execution-Guided Retrieval and Reuse Framework for Reliable Machine Learning Pipeline Code Generation

Wenyu Zhao, Tingjie Chen, Boyuan Wang, Zijian Shen
pp.286-297 DOI: 10.62051/46qs2n55 August 14, 2026
Download PDF
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.
Keywords: Retrieval-augmented generation; Code reuse; Machine-learning pipelines; Execution feedback; Program repair; Scikit-learn.
APA Citation: Wenyu Zhao, Tingjie Chen, Boyuan Wang, Zijian Shen (2026). PIPE-RAG: An Execution-Guided Retrieval and Reuse Framework for Reliable Machine Learning Pipeline Code Generation. Transactions on Computer Science and Intelligent Systems Research, 13(1), 286-297. https://doi.org/10.62051/46qs2n55

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