Query2Diagram: Answering Developer Queries with UML Diagrams

1. Problem Statement

Software documentation, particularly UML diagrams, is often outdated or non-existent, creating a significant barrier to understanding complex codebases. While automated reverse engineering (RE) tools can generate UML diagrams from code, they typically produce exhaustive structural dumps containing overwhelming detail that hinders comprehension. These tools lack semantic context and developer intent, failing to filter for relevant components.

Conversely, while Large Language Models (LLMs) excel at understanding code and natural language, their application to query-driven diagram generation (creating diagrams that answer specific questions about code) remains unexplored. Existing LLM approaches either generate complete diagrams from code (ignoring user intent) or derive diagrams from textual specifications (not code). There is a critical gap in generating focused, contextual UML views directly from source code based on natural language queries.

2. Methodology

The authors propose a novel pipeline called Query2Diagram, which generates UML diagrams in response to specific developer queries. The methodology consists of three main stages:

A. Data Curation

• Code Selection: The team collected the 150 most-starred open-source repositories (MIT/Apache licenses) covering 12 languages (C, C++, Java, Python, JS, TS, Rust, etc.). They filtered for files between 3K–15K characters to ensure manageable complexity.
• Query Synthesis: Since no public benchmark links code queries to UML diagrams, they used two open LLMs (DeepSeek-R1 and QwQ) to generate natural language questions focusing on architecture, API usage, and design patterns.
• Diagram Construction (Ground Truth):
  • Initial diagrams were generated using Claude 3.5 Sonnet to ensure high JSON validity.
  • The output format is a structured JSON graph containing lists of nodes, edges, and packages.
  • Node Types: Class, method, field, function, variable, and abstract entity.
  • Correction: All generated graphs were manually reviewed and corrected to fix defects (e.g., hallucinated nodes, syntax errors, missing connections). Irreparable graphs were discarded.
  • Result: A curated dataset of code-query-diagram triples (264 training, 36 validation).

B. Model Adaptation (Fine-Tuning)

• Base Model: Qwen2.5-Coder-14B-Instruct was selected for its strong zero-shot coding performance and compatibility with single-GPU training (NVIDIA Tesla V100).
• Technique: QLoRA (Quantized Low-Rank Adaptation) was used to fine-tune the model.
  • Weights were quantized to 4-bit (NF4).
  • Full-precision adapters were trained using supervised cross-entropy loss.
  • Optimizations included Unsloth, L2 regularization, dropout, and gradient accumulation.
• Training Objective: The model learns to map a (code file, natural language query) pair to a structured JSON representation of a UML diagram.

C. Inference and Generation

• Constraint: To ensure syntactic validity, the model uses Outlines to impose a JSON schema as a constrained grammar during decoding.
• Output: The model generates three versions of the diagram (Minimal, Medium, Full) and a text summary, allowing users to choose the level of abstraction.

3. Key Contributions

1. Task Formulation: Introduced the novel task of query-driven UML generation, where diagrams are not exhaustive but are filtered and abstracted based on specific developer information needs.
2. Structured Generation Method: Developed a practical pipeline using fine-tuned LLMs to output JSON-based intermediate representations (nodes, edges, packages) rather than raw markup (PlantUML/Mermaid), significantly reducing syntax errors.
3. Comprehensive Evaluation Framework: Created a dual-evaluation strategy combining:
   • Automatic Defect Analysis: Scanning for 19 defect patterns (stylistic, structural, unrenderable).
   • Human Relevance Assessment: Classifying nodes into four categories: Sufficiency (essential), Completeness (helpful), Hallucination (invented), and Verbosity (irrelevant).
4. Empirical Evidence: Demonstrated that fine-tuning on a modest amount of manually corrected data yields dramatic improvements over both base LLMs and synthetic-only training.

4. Results

The study benchmarked five models: GPT-4o, Claude 3.5 Sonnet, base Qwen2.5-Coder-14B, and two fine-tuned versions (trained on synthetic data vs. manually corrected data).

• Defect Reduction: The fine-tuned model trained on manually corrected data (Qwen2.5-Coder-14B SFT: Fixed Claude Synth) achieved the lowest defect rates, outperforming GPT-4o and Claude 3.5 Sonnet. It eliminated "unacceptable" defects entirely.
• Semantic Relevance (F1 Scores):
  • Base Models: Showed a trade-off. GPT-4o/Claude had high precision (few hallucinations) but low recall (missed relevant info). Base Qwen had high recall but low precision (many hallucinations/verbosity).
  • Fine-Tuned Model: Successfully closed this gap. It achieved the highest F1 scores (Micro: 0.808, Macro: 0.823) by retaining the recall of the base Qwen while drastically suppressing hallucinations and verbosity.
• Hard Metrics: Even when evaluating only "Sufficiency" nodes (indispensable elements), the fine-tuned model maintained high precision and recall, proving it generates structurally sound and semantically faithful diagrams.

5. Significance and Future Work

• Feasibility of On-Demand Documentation: The paper proves that LLMs can generate scalable, contextual documentation that adapts to developer queries, moving beyond static, outdated diagrams.
• Data Efficiency: The results highlight that high-quality, small-scale datasets (manually corrected) are more effective for specialized tasks than massive, uncurated synthetic data.
• Future Directions:
  • Extending from single-file analysis to multi-file/project-wide dependency mapping.
  • Developing conversational interfaces for iterative refinement (e.g., "zoom into error handling").
  • Expanding to other diagram types (Sequence, Activity, State Machine).
  • Integrating with IDEs for real-time, zero-effort documentation.

In conclusion, Query2Diagram establishes a new paradigm for software engineering tools, demonstrating that combining LLMs with structured output constraints and targeted fine-tuning can effectively bridge the gap between code complexity and developer understanding.