Embedding Gemma-2
for Washington State Law

US-State-Law-EmbeddingGemma-2

This is a highly specialized embedding model fine-tuned specifically for the Revised Code of Washington (RCW). Designed to run on consumer hardware such as mobile devices and laptops, delivering low-latency semantic representations for on-device applications, like search, retrieval-augmented generation (RAG), classification, and clustering.

504.1k

RCW context pairs

768

Vector Dimension

87.9%

Accuracy@10

+49.6%

MRR Improvement

Hugging Face

Download ›

Performance Analysis

This section details the critical retrieval metrics comparing the original baseline model against our fine-tuned custom model. Evaluated against a held-out validation set of ~47,000 unique Washington State law queries, the fine-tuning process resulted in massive improvements in recall and ranking accuracy, proving its efficacy as the backbone for legal RAG pipelines.

Retrieval Metrics Comparison

87.89%

Fine-Tuned Accuracy@10

80.27%

Fine-Tuned Accuracy@5

57.09%

Fine-Tuned MRR@10

64.64%

Fine-Tuned NDCG@10

Outperforms Larger Domain-Specific Models

The US-State-Law-EmbeddingGemma-2 (270M) model establishes a new performance frontier on the Washington Law Benchmark (WLB) by achieving the highest retrieval accuracy while optimizing parameter efficiency. Our fine-tuned 270M parameter architecture achieves an optimal accuracy-to-scale ratio, securing a ~24 percentage point gain over the standard EmbeddingGemma 2 baseline and outperforming larger, domain-specific models like WA-State-Law-Large (335M).

The Washington law Benchmark (WLB)

Retrieval Accuracy vs. Model Size on the WLB Benchmark

Implementation Guide

Integrating this custom model into your Python stack is straightforward via the sentence-transformers library. However, strict adherence to the query formatting rules below is absolutely critical to achieve the benchmarked accuracy.

⚠ 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.

Dependencies

pip install -U sentence-transformers

Python Usage Example

import torch
from sentence_transformers import SentenceTransformer, util

# 1. Load the 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
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)

# 7. 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}")

Embedding Gemma-2 for Washington State Law