Extracting Logical Structure in Code Documents via Semantic Segmentation and Language Models

Abstract

A key task for better understanding and maintaining source code as a code document is semantic code document segmentation, i. e., dividing the code into coherent blocks of functional intent. Despite its potential to enhance code comprehension, navigation, and reuse, this task remains underexplored due to the lack of semantically segmented datasets. Prior work generally relied on syntactic signals (e. g., AST-based features) or manual heuristics, limiting scalability and generalization. We propose two language-model-based strategies: (i) a line-by-line approach that classifies each line of code separately before grouping the results into functional units, and (ii) a range-based approach that aims to directly determine groups of code lines from the input. The latter is particularly suitable for generative language models as they can take an entire code file as the input context. Furthermore, we release two expert-annotated datasets from real-world scientific code in both a low-resource language, R, and the widely used Python language. Experiments show that line-by-line strategy with a local context of K surrounding lines generally outperforms the range-based approach for both programming languages. Fine-tuning smaller models like CodeBERT and CodeT5+ for line-by-line classification generally outperforms larger, generative language models, even without R-specific pretraining. On a single GPU, the runtime of CodeBERT is 100-170x faster than those of the best competing LLMs, supporting the practical integration of semantic segmentation into modern development environments. The code, prompts, and datasets are available at: https://github.com/Dahouabdelhalim/CodeSeg