Query2Diagram: Answering Developer Queries with UML Diagrams

TL;DR

Query2Diagram leverages fine-tuned LLMs to generate UML diagrams that directly answer developer questions, improving relevance and correctness over traditional reverse engineering tools.

Contribution

It introduces a query-driven UML diagram generation method using fine-tuned LLMs, focusing on semantic relevance and structural correctness based on developer queries.

Findings

Fine-tuning on curated data significantly improves diagram quality.

The approach achieves higher F1 scores than existing LLMs.

Generated diagrams are both structurally sound and semantically faithful.

Abstract

Software documentation frequently becomes outdated or fails to exist entirely, yet developers need focused views of their codebase to understand complex systems. While automated reverse engineering tools can generate UML diagrams from code, they produce overwhelming detail without considering developer intent. We introduce query-driven UML diagram generation, where LLMs create diagrams that directly answer natural language questions about code. Unlike existing methods, our approach produces semantically focused diagrams containing only relevant elements with contextual descriptions. We fine-tune Qwen2.5-Coder-14B on a curated dataset of code files, developer queries, and corresponding diagram representations in a structured JSON format, evaluating with both automatic detection of structural defects and human assessment of semantic relevance. Results demonstrate that fine-tuning on a modest amount of manually corrected data yields dramatic improvements: our best model achieves the highest F1 scores while reducing defect rates below state-of-the-art LLMs, generating diagrams that are both structurally sound and semantically faithful to developer queries. Thus, we establish the feasibility of using LLMs for scalable contextual, on-demand documentation generation. We make our code and dataset publicly available at https://github.com/i-need-a-pencil/query2diagram.