Hugging Face | GitHub | Launch Blog | Documentation |
License: Apache 2.0

US-State-Law-EmbeddingGemma-2-GGUF

US-State-Law-EmbeddingGemma-2-GGUF is a highly efficient, 4-bit quantized(q4_k_m) version of our specialized embedding model fine-tuned specifically for Legal Information Retrieval (IR) within the State of Washington. Designed to run on consumer hardware such as mobile devices and laptops, US-State-Law-EmbeddingGemma-2-GGUF delivers low-latency semantic representations for on-device applications, like search, retrieval-augmented generation (RAG), classification, and clustering.

Generic embedding models often perform suboptimally on legal texts due to the semantic gap between natural language questions (e.g., "What dollar amount makes a theft a first degree felony?") and formal statutory legalese. This model bridges that gap, allowing plain-English queries, legal scenarios, and document drafts to be accurately mapped to their corresponding Washington State statutes (Revised Code of Washington - RCW).

This specific version delivers the powerful retrieval capabilities of the LoRA fine-tuned EmbeddingGemma 2 architecture in a lightweight GGUF format, making it ideal for low-resource environments, CPU-only deployments, and edge computing without sacrificing significant retrieval accuracy.

EmbeddingGemma 2 is an open multimodal embedding model built by Google DeepMind which maps text (incl. code), images, video, and audio inputs—and combinations thereof—into a single, unified 768-dimensional vector space. The model has 740M total parameters, combining a 270M parameter text model with modular vision (170M) and audio (300M) encoders.

Available Models

Model Overview

• Base Model: EmbeddingGemma 2
• Task: Semantic Search / Information Retrieval / Legal Preemption Analysis
• Language: English (Legal Domain)
• Format/Precision: GGUF / 4-bit(q4_k_m)
• Similarity Function: Cosine Similarity

Key Features

• Low Memory Footprint: 4-bit quantization drastically reduces RAM/VRAM requirements for hosting.
• Edge & CPU Optimized: The GGUF format enables fast, efficient inference on standard consumer hardware and CPU-only servers.
• Domain Specific: Retains the high-quality legal retrieval performance fine-tuned specifically for the Washington State legal domain (RCW).
• Plain-English Queries: Supports complex, natural language legal questions mapped directly to statute

Intended Use Cases

This model is optimized to act as the retriever component in legal Retrieval-Augmented Generation (RAG) pipelines. Primary use cases include:

1. Statutory Cross-Referencing: Mapping natural language legal questions to specific RCWs.
2. Preemption Checking: Automatically retrieving state laws that may preempt or conflict with proposed municipal ordinances.
3. Legal Research Automation: Clustering and searching local agency drafts against established state frameworks.
4. AI Legal Assistants: Powering chatbots and research tools that require accurate retrieval of Washington State laws before generating an answer.
5. Automated Compliance: Scanning contracts or external drafts against established state legislative frameworks.

Technical Details & Training Methodology

The Semantic Gap

A standard dense retriever often fails on legal tasks because it relies on vocabulary overlap rather than conceptual legal mapping. To address this, US-state-law-EmbeddingGemma-2-GGUF was fine-tuned using a synthetic, high-variance dataset designed to map hypothetical language to actual statutes.

Training Data

The model was fine-tuned on synthetic legal query–passage pairs generated from Washington State RCW statutes.

The dataset spans hundreds of legal categories derived from the RCW structure and includes:

• Natural language paraphrases of legal questions
• Hypothetical legal scenarios
• Statute-grounded positive document matches

Hyperparameters & Architecture

• Fine-Tuning Method: LoRA (Low-Rank Adaptation)
• Precision: 16-bit (fp16)
• Loss Function: Multiple Negatives Ranking (MNR) Loss
• Batch Size: 128
• Epochs: 1
• batch_sampler: no_duplicates
• multi_dataset_batch_sampler: round_robin
• Learning Rate Decay: Linear
• Infrastructure: High-Performance Computing (HPC) Cluster

All Hyperparameters

Click to expand

• per_device_train_batch_size: 128
• num_train_epochs: 3.0
• max_steps: 1780
• learning_rate: 2e-05
• lr_scheduler_type: linear
• lr_scheduler_kwargs: None
• warmup_steps: 0.1
• optim: adamw_torch_fused
• optim_args: None
• weight_decay: 0.0
• adam_beta1: 0.9
• adam_beta2: 0.999
• adam_epsilon: 1e-08
• optim_target_modules: None
• gradient_accumulation_steps: 2
• average_tokens_across_devices: True
• max_grad_norm: 1.0
• label_smoothing_factor: 0.0
• bf16: True
• fp16: False
• bf16_full_eval: False
• fp16_full_eval: False
• tf32: None
• gradient_checkpointing: False
• gradient_checkpointing_kwargs: None
• torch_compile: False
• torch_compile_backend: None
• torch_compile_mode: None
• use_liger_kernel: False
• liger_kernel_config: None
• use_cache: False
• neftune_noise_alpha: None
• torch_empty_cache_steps: None
• auto_find_batch_size: False
• log_on_each_node: True
• logging_nan_inf_filter: True
• include_num_input_tokens_seen: no
• log_level: passive
• log_level_replica: warning
• disable_tqdm: False
• project: huggingface
• trackio_space_id: None
• trackio_bucket_id: None
• trackio_static_space_id: None
• per_device_eval_batch_size: 128
• prediction_loss_only: True
• eval_on_start: False
• eval_do_concat_batches: True
• eval_use_gather_object: False
• eval_accumulation_steps: None
• include_for_metrics: []
• batch_eval_metrics: False
• save_only_model: False
• save_on_each_node: False
• enable_jit_checkpoint: False
• push_to_hub: False
• hub_private_repo: None
• hub_model_id: None
• hub_strategy: every_save
• hub_always_push: False
• hub_revision: None
• load_best_model_at_end: True
• ignore_data_skip: False
• restore_callback_states_from_checkpoint: False
• full_determinism: False
• seed: 42
• data_seed: None
• use_cpu: False
• accelerator_config: {'split_batches': False, 'dispatch_batches': None, 'even_batches': True, 'use_seedable_sampler': True, 'non_blocking': False, 'gradient_accumulation_kwargs': None}
• parallelism_config: None
• dataloader_drop_last: False
• dataloader_num_workers: 0
• dataloader_pin_memory: True
• dataloader_persistent_workers: False
• dataloader_prefetch_factor: None
• dataloader_multiprocessing_context: None
• dataloader_in_order: True
• remove_unused_columns: True
• label_names: None
• train_sampling_strategy: random
• length_column_name: length
• ddp_find_unused_parameters: None
• ddp_bucket_cap_mb: None
• ddp_broadcast_buffers: False
• ddp_static_graph: None
• ddp_backend: None
• ddp_timeout: 1800
• fsdp: None
• fsdp_config: None
• deepspeed: None
• debug: []
• skip_memory_metrics: True
• do_predict: False
• resume_from_checkpoint: None
• local_rank: -1
• prompts: {'question': 'task: search result | query: ', 'passage_text': 'title: none | text: '}
• batch_sampler: no_duplicates
• multi_dataset_batch_sampler: proportional
• router_mapping: {}
• learning_rate_mapping: {}
• warmup_ratio: None

Training Logs

Click to expand

• The bold row denotes the saved checkpoint.

Training Time

• Training: 8.3 hours

Evaluation Metrics

The model was evaluated on a rigorously held-out validation set of synthetic municipal drafts mapped 1-to-1 against Washington State RCWs. The metrics below demonstrate the model's exceptional retrieval capabilities.

Interpretation: When a user asks this model a legal question in plain English, there is an 87.9% probability that the exact governing state law will be returned in the top 10 search results.

Limitations

• This model does not provide legal advice.
• Performance is limited to Washington State law (RCW) and may not generalize to other jurisdictions.
• Outputs depend on the quality of the underlying document corpus.
• Should be used as a retrieval tool, not a final decision-making system.

Usage Examples

Semantic Search with llama-cpp-python

To run the GGUF quantized model, you will need to use a framework that supports GGUF embeddings, such as llama-cpp-python.

import numpy as np
from llama_cpp import Llama

# 1. Load the quantized GGUF model
# Make sure to set embedding=True
model = Llama(
    model_path="./US-state-law-EmbeddingGemma-2-q4_k_m.gguf",
    embedding=True,
    verbose=False
)

# 2. Define the laws (Your Vector Database)
laws = [
    "RCW 9A.56.030: Theft in the first degree. A person is guilty of theft in the first degree if he or she commits theft of property or services which exceed(s) five thousand dollars in value.",
    "RCW 46.61.502: Driving under the influence. A person is guilty of driving while under the influence of intoxicating liquor...",
    "RCW 9A.36.011: Assault in the first degree. A person is guilty of assault in the first degree if he or she..."
]

# 3. Define the user's search query
user_query = "What dollar amount makes a theft a first degree felony?"

# 4. Generate embeddings
# llama-cpp-python returns embeddings as lists of floats
query_embedding = np.array(model.embed(user_query))
law_embeddings = [np.array(model.embed(law)) for law in laws]

# 5. Calculate Cosine Similarity
def cosine_similarity(v1, v2):
    return np.dot(v1, v2) / (np.linalg.norm(v1) * np.linalg.norm(v2))

cosine_scores = [cosine_similarity(query_embedding, doc) for doc in law_embeddings]

# 6. Print the top result
best_idx = np.argmax(cosine_scores)
print(f"Top Match: {laws[best_idx]}")
print(f"Similarity Score: {cosine_scores[best_idx]:.4f}")

Best Practices

For optimal embedding quality and runtime efficiency, follow these configurations and best practices:

1. Task Instruction Prefixes

EmbeddingGemma 2 is trained with short task instruction prefixes prepended to text inputs. Using the right prefix improves quality; omitting it still works but reduces precision. Prefixes apply to text only. Pass images, video, and audio without any prefix.

Documents with a real title should be formatted as title: {title} | text: {content}. Use title: none when no title is available.

Prefix Notation & Usage

We offer two types of task prefixes, depending on how embeddings are used in the task. There are two types of tasks:

• Asymmetric Tasks (e.g. retrieval): Use a query prefix for queries and a document prefix for corpus items.
• Symmetric Tasks (e.g. classification, similarity): Apply the same task prefix to all inputs being compared.

Please note: prompt_name="Document" applies title: none; titled documents must still be formatted manually, like model.encode(f"title: {title} | text: {document}")

Model Citation

@misc{US_state_law_embedding_gemma2_2026,
  title={US-state-law-EmbeddingGemma-2-GGUF: Fine-Tuned Dense Retrieval for Washington State La},
  author={Tomar, Shlok and Complex Social Interactions Lab},
  year={2026},
  publisher={Hugging Face},
  howpublished={\url{https://huggingface.co/CSI-lab/US-state-law-EmbeddingGemma-2-GGUF}},
  note={Hugging Face Model Repository}
}

BibTeX

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084",
}

Low-Rank Adaptation (LoRA)

@article{hu2021lora,
  title={LoRA: Low-Rank Adaptation of Large Language Models},
  author={Hu, Edward J and Shen, Yelong and Wallis, Phil and Allen-Zhu, Zeyuan and Li, Yuanzhi and Wang, Shean and Wang, Lu and Chen, Weizhu},
  journal={arXiv preprint arXiv:2106.09685},
  year={2021}
}

EmbeddingGemma Paper

@article{embedding_gemma_2025,
  title={EmbeddingGemma: Powerful and Lightweight Text Representations},
  author={Schechter Vera, Henrique and Dua, Sahil and Team, EmbeddingGemma},
  journal={arXiv preprint arXiv:2509.20354},
  year={2025},
  url={https://arxiv.org/abs/2509.20354}
}

EmbeddingGemma 2 Official Release

@online{embedding_gemma_2_2026,
  author={Google DeepMind},
  title={EmbeddingGemma 2 is a best-in-class open model for natively multimodal embeddings},
  year={2026},
  url={https://blog.google/innovation-and-ai/technology/developers-tools/embeddinggemma-2/},
  note={Google Innovation and AI Blog}
}