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

US-State-Law-EmbeddingGemma-2

US-State-Law-EmbeddingGemma-2 is a highly 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 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 version leverages the EmbeddingGemma 2 architecture and was fine-tuned using Low-Rank Adaptation (LoRA) before being merged into a full 16-bit precision model for optimal performance and efficiency.

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)
• Precision: Full 16-bit (Merged LoRA)
• Similarity Function: Cosine Similarity

Key Features

• Fine-tuned specifically for the Washington State legal domain (RCW).
• Parameter-efficient fine-tuning via LoRA, yielding high performance with reduced training overhead.
• Merged into a full 16-bit model for simplified deployment.
• Supports complex, natural language legal queries.
• Designed for RAG-based legal assistants and semantic search pipelines.

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 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 sentence-transformers

import torch
from sentence_transformers import SentenceTransformer, util

# 1. Load the 16-bit fine-tuned model
model = SentenceTransformer('CSI-lab/US-state-law-EmbeddingGemma-2', 
 model_kwargs={
        "torch_dtype": torch.bfloat16,
        "config_kwargs": {
            "load_vision_encoder": False,
            "load_audio_encoder": 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. Encode the documents and the query with respective prompt_name parameters
query_embedding = model.encode(user_query, convert_to_tensor=True, prompt_name="SearchQuery")
law_embeddings = model.encode(laws, convert_to_tensor=True, prompt_name="Document")

# 5. Calculate Cosine Similarity
cosine_scores = util.cos_sim(query_embedding, law_embeddings)

# 6. Print the top result
best_idx = cosine_scores.argmax().item()
print(f"Top Match: {laws[best_idx]}")
print(f"Similarity Score: {cosine_scores[0][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: 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}},
  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}
}