gemma-4-sql

Natural text to SQL with Gemma 4; with DuckDB support and swappable-backends: PyTorch (HF); PyTorch (Native); Keras ; JAX; JAX / MaxText.

Documentation:

• Extending / Custom Backends
• Deployment & CI/CD
• DuckDB Support
• SQL Dataset Analysis
• Architecture Details
• Usage Guide

gemma-4-sql is a specialized SDK and CLI tool designed for orchestrating Text-to-SQL training pipelines. It provides an end-to-end framework capable of ingesting diverse Text-to-SQL datasets, transforming them using Google's grain library into consistent multidimensional formats, and preparing them for modern AI-Hypercomputer workloads.

We explicitly integrate with and support the following Gemma 4 model architectures across different ecosystems:

• PyTorch (HF): Directly imports and uses Gemma4ForCausalLM from Hugging Face Transformers;
• PyTorch (Native): A custom, from-scratch implementation of Gemma 4 built with Native PyTorch (torch.nn);
• MaxText: Directly imports and uses Gemma4Model from AI-Hypercomputer MaxText;
• JAX: A custom, from-scratch implementation of Gemma 4 built with Flax NNX;
• Keras: Directly imports and uses GemmaCausalLM from KerasNLP;

Feature Support Matrix

Note on ETL differences: JAX, MaxText, and Keras all leverage Google's grain library. While JAX and Keras use a shared BaseFormatTransform yielding standard inputs and targets, MaxText uses MaxTextFormatTransform to inject additional Seq2Seq features like segment_ids and positions expected by the MaxText architecture. Distributed environments use JAXDistributedSharding.

Documentation & Usage

For full instructions on Installation, Development, ETL, Training, Inference, and other workflows, please see the Usage Guide.

For a comprehensive guide on running these training scripts across distributed infrastructure (like Google Cloud TPU VMs) using MaxText or JAX, please refer to the DEPLOY_TO_TPU.md file.

---

License

Licensed under either of

• Apache License, Version 2.0 (LICENSE-APACHE or https://www.apache.org/licenses/LICENSE-2.0)
• MIT license (LICENSE-MIT or https://opensource.org/licenses/MIT)

at your option.

Contribution

Unless you explicitly state otherwise, any contribution intentionally submitted for inclusion in the work by you, as defined in the Apache-2.0 license, shall be dual licensed as above, without any additional terms or conditions.