LRriver opened a new pull request, #352: URL: https://github.com/apache/hugegraph-ai/pull/352
## Summary This PR adds a new `text2gremlin/AST_Text2Gremlin` module for Text2Gremlin data generation, LLM-based data augmentation, syntax validation, dataset merging, and DPO data construction. The goal is to provide a reproducible data pipeline for generating high-quality Gremlin query training data from structured graph schemas and templates. The module starts from schema-aware Gremlin template generation, then uses LLMs to produce natural-language instructions, migrates samples across graph scenarios, validates generated Gremlin syntax, merges SFT data, and optionally builds DPO preference data by comparing correct Gremlin queries with Groovy-style negative samples. The generated Text2Gremlin dataset has also been published on Hugging Face: https://huggingface.co/datasets/Lriver/Text2Gremlin ## Motivation Text2Gremlin needs training data that is both syntactically valid and diverse enough to cover common graph query patterns. Manually writing these samples is expensive and difficult to scale, especially when the data needs to cover different graph domains, operation types, traversal structures, and natural-language styles. This module provides a structured generation and augmentation pipeline: 1. Generate Gremlin queries from schema-aware templates. 2. Validate query syntax with an ANTLR-based Gremlin parser. 3. Translate Gremlin queries into multiple natural-language instruction styles. 4. Migrate existing samples to other graph scenarios. 5. Merge and deduplicate SFT-style training data. 6. Generate DPO preference data for Groovy-vs-Gremlin alignment. ## Project structure This PR adds the following Text2Gremlin module structure: ```text text2gremlin/AST_Text2Gremlin/ ├── README.md ├── README_zh.md ├── requirements.txt ├── config_example.json ├── generate_corpus.py ├── analyze_syntax.py ├── run_llm_pipeline.py ├── gremlin_templates.csv ├── base/ │ ├── Config.py │ ├── Schema.py │ ├── GremlinBase.py │ ├── GremlinExpr.py │ ├── GremlinParse.py │ ├── GremlinTransVisitor.py │ ├── TraversalGenerator.py │ ├── CombinationController.py │ ├── generator.py │ ├── combination_control_config.json │ ├── gremlin/ │ │ ├── Gremlin.g4 │ │ ├── GremlinLexer.py │ │ ├── GremlinParser.py │ │ ├── GremlinVisitor.py │ │ └── GremlinListener.py │ └── template/ │ ├── schema_dict.txt │ └── syn_dict.txt ├── db_data/ │ ├── schema/ │ │ └── movie_schema.json │ ├── movie/raw_data/ │ │ ├── vertex_*.csv │ │ └── edge_*.csv │ └── reference/ │ └── schemas_data.json ├── llm_augment/ │ ├── generalize_llm.py │ ├── migrate_scenario.py │ ├── merge_dataset.py │ └── generate_dpo_data.py └── tests/ ├── test_analyze_syntax.py ├── test_generate_corpus.py ├── test_gremlin_base.py ├── test_generalize_llm.py ├── test_migrate_scenario.py ├── test_merge_dataset.py └── test_run_llm_pipeline.py ``` The module is organized around four layers: - `base/`: schema-aware Gremlin AST parsing, traversal generation, recipe representation, grammar integration, and generation controls. - `db_data/`: seed graph schema and movie-domain raw data used by the generator. - `llm_augment/`: LLM-based translation, scenario migration, dataset merge, and DPO data construction. - `tests/`: focused tests for generation, augmentation, merge behavior, pipeline argument forwarding, and syntax analysis. ## What changed ### 1. AST-based Gremlin corpus generation This PR adds a schema-aware Gremlin generation framework under `text2gremlin/AST_Text2Gremlin/base`. Key pieces include: - ANTLR Gremlin grammar/parser/visitor integration. - Recipe-style Gremlin representation for query construction. - Step, predicate, anonymous traversal, connector, and terminal handling. - Recursive traversal generation with schema constraints. - Connectivity validation for generated paths. - Data value filling from schema and raw graph data. - Deduplication and syntax filtering for generated queries. - Combination control through `base/combination_control_config.json`. The main entry point is: ```bash python generate_corpus.py ``` The generation process can use: - `db_data/schema/movie_schema.json` - `db_data/movie/raw_data/*.csv` - `gremlin_templates.csv` - `base/template/schema_dict.txt` - `base/template/syn_dict.txt` ### 2. Movie-domain seed schema and data This PR includes a movie graph scenario used as the initial schema/data source for Text2Gremlin generation. Added files include: - `db_data/schema/movie_schema.json` - movie vertex CSV files, such as `vertex_movie.csv`, `vertex_person.csv`, `vertex_user.csv` - movie edge CSV files, such as `edge_acted_in.csv`, `edge_directed.csv`, `edge_rate.csv` - `db_data/reference/schemas_data.json` These files provide the schema, labels, properties, and example values needed by the generation pipeline. ### 3. LLM-based natural-language augmentation This PR adds `llm_augment/generalize_llm.py`, which converts generated Gremlin queries into natural-language Text2Gremlin samples. The augmentation supports multiple instruction styles, including: - direct question style - command style - natural conversational style - domain-aware wording The output is designed for SFT-style Text2Gremlin training data, where each sample pairs a user instruction with a valid Gremlin query. ### 4. Scenario migration augmentation This PR adds `llm_augment/migrate_scenario.py`, which migrates generated samples from one graph scenario to another. The default migration mode is now: ```text same_operation ``` In this mode, the model is asked to migrate a source query into the target scenario while preserving the original operation type, such as read, create, update, or delete. It can generate multiple target-scenario samples for each source sample. The previous broader behavior is still available through: ```text mixed_operations ``` In this mode, the model may generate target-scenario samples across multiple operation types. This is useful when users want more diverse CRUD-style augmentation, but it is no longer the default because not every source query is suitable for migration into every operation type. The number of same-operation migration samples is configurable and defaults to `3`. ### 5. Dataset merge and statistics This PR adds `llm_augment/merge_dataset.py` for merging augmented Text2Gremlin data into final SFT-style outputs. The merge step supports: - combining direct translation data and scenario migration data - filtering invalid or incomplete samples - deduplicating samples - collecting domain and CRUD operation statistics - exporting train/validation style data files This keeps generated data preparation separate from model training, making it easier to inspect and reuse the dataset. ### 6. DPO data generation This PR adds `llm_augment/generate_dpo_data.py` for generating DPO preference data. The DPO generation compares valid Gremlin answers with lower-quality Groovy-style or non-preferred outputs. This is intended to help align models toward producing Gremlin queries instead of code-like alternatives. ### 7. End-to-end pipeline runner This PR adds `run_llm_pipeline.py` as a staged pipeline entry point. The pipeline supports running stages such as: ```text translate -> migrate -> merge -> dpo ``` This makes it possible to run only the required part of the pipeline during development or data refresh. ### 8. Syntax analysis tooling This PR adds `analyze_syntax.py` for analyzing generated Gremlin query distributions. The analysis can report: - Gremlin step frequency - predicate usage - traversal pattern distribution - operation-type distribution - syntax coverage statistics This is useful for checking whether generated data is overly concentrated on a small set of Gremlin patterns. ### 9. Configuration and examples This PR adds: - `config_example.json` - `requirements.txt` - English README - Chinese README The config example documents model API settings and generation-related options. Sensitive local config files are excluded by `.gitignore`. The README files describe: - module purpose - installation - configuration - corpus generation - LLM augmentation - scenario migration modes - dataset merge - DPO generation - Hugging Face dataset location - expected output files ### 10. Tests This PR adds focused pytest coverage under `text2gremlin/AST_Text2Gremlin/tests`. The tests cover: - Gremlin syntax analysis - corpus generation behavior - Gremlin base parsing/generation helpers - LLM generalization helpers - scenario migration modes - same-operation filtering - mixed-operation mode behavior - merge dataset behavior - pipeline argument forwarding - config loading - dictionary fallback and dictionary loading behavior ## Dataset The generated dataset is available here: https://huggingface.co/datasets/Lriver/Text2Gremlin The dataset is intentionally hosted outside this repository to avoid committing large generated artifacts. This repository contains the generation code, schema/data seeds, configuration examples, and documentation needed to reproduce or extend the dataset. ## Compatibility and scope This PR only adds the new Text2Gremlin generation module under: ```text text2gremlin/AST_Text2Gremlin ``` It does not change the existing HugeGraph LLM runtime APIs, HugeGraph Python client APIs, or other existing modules. Large generated output artifacts are not committed. Local model API configuration is expected to be stored in a local config file and is ignored by git. ## Validation The same head branch has already been validated through the existing PR workflow in `hugegraph/hugegraph-ai#52`, including Ruff checks, dependency/license checks, module CI checks, and CodeRabbit review. Local validation used during preparation: ```bash uv run --with-requirements text2gremlin/AST_Text2Gremlin/requirements.txt pytest text2gremlin/AST_Text2Gremlin/tests -q uv run ruff format --check . uv run ruff check . git diff --check ``` ## Related This PR is based on the same `text2gremlin` branch as: https://github.com/hugegraph/hugegraph-ai/pull/52 -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
