AstraAI用LLM+检索+语法树分析，帮科学家在超算代码中精准生成符合规范的代码。
AstraAI: LLMs, Retrieval, and AST-Guided Assistance for HPC Codebases
- 结合检索和语法树分析，生成带上下文的精准代码提示。
- 在AMReX项目上实现结构一致的代码修改，保持原有编程风格。
- 支持本地与云端模型，适合超算科研人员高效开发复杂科学代码。
我们提出AstraAI，一种面向高性能计算（HPC）软件开发的命令行编码框架。该框架直接运行于Linux终端，集成大语言模型（LLM）、检索增强生成（RAG）和抽象语法树（AST）结构分析，实现对复杂科学代码库的上下文感知代码生成。核心思路是构建高保真提示，将用户请求与通过RAG从底层代码库检索的相关代码片段、以及来自AST分析提取的结构上下文相结合，使模型能准确理解相关函数、数据结构及整体代码组织。框架支持在限定范围内修改源码，同时保持与周围代码的结构一致性。AstraAI支持从Hugging Face部署的本地模型和通过美国科学云访问的API前沿模型，可在不同HPC环境中灵活部署。系统生成的代码与现有项目结构和编程模式高度一致。我们在AMReX——一个由美国能源部支持的用于百亿亿次计算应用的HPC软件基础设施——中展示了AstraAI在典型HPC代码生成任务上的有效性。
原文摘要 · Abstract (English)
We present AstraAI, a command-line interface (CLI) coding framework for high-performance computing (HPC) software development. AstraAI operates directly within a Linux terminal and integrates large language models (LLMs) with Retrieval-Augmented Generation (RAG) and Abstract Syntax Tree (AST)-based structural analysis to enable context-aware code generation for complex scientific codebases. The central idea is to construct a high-fidelity prompt that is passed to the LLM for inference. This prompt augments the user request with relevant code snippets retrieved from the underlying framework codebase via RAG and structural context extracted from AST analysis, providing the model with precise information about relevant functions, data structures, and overall code organization. The framework is designed to perform scoped modifications to source code while preserving structural consistency with the surrounding code. AstraAI supports both locally hosted models from Hugging Face and API-based frontier models accessible via the American Science Cloud, enabling flexible deployment across HPC environments. The system generates code that aligns with existing project structures and programming patterns. We demonstrate AstraAI on representative HPC code generation tasks within AMReX, a DOE-supported HPC software infrastructure for exascale applications.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。