import torch
from sentence_transformers import SentenceTransformer


similarities_per_attn = []
rankings_per_attn = []
for attention in ["eager", "sdpa", "flash_attention_2"]:
    model = SentenceTransformer(
        "google/embeddinggemma-300m", 
        model_kwargs={"attn_implementation": f"{attention}", "dtype": torch.bfloat16}
    )

    query = "Which planet is known as the Red Planet?" * 30
    documents = [
        "Venus is often called Earth's twin because of its similar size and proximity." * 30,
        "Mars, known for its reddish appearance, is often referred to as the Red Planet." * 30,
        "Jupiter, the largest planet in our solar system, has a prominent red spot." * 30,
        "Saturn, famous for its rings, is sometimes mistaken for the Red Planet." * 30
    ]
    query_embeddings = model.encode_query(query)
    document_embeddings = model.encode_document(documents)

    # Compute similarities to determine a ranking
    similarities = model.similarity(query_embeddings, document_embeddings)
    similarities_per_attn.append(similarities)

    # Convert similarities to a ranking
    ranking = similarities.argsort(descending=True)[0]
    rankings_per_attn.append(ranking)

    print(f"{attention}: ")
    print(similarities)
    print(ranking)
    print()


similarity, rank = similarities_per_attn[0], rankings_per_attn[0]
for i in range(1, len(similarities_per_attn)):
    if not (torch.allclose(similarity, similarities_per_attn[i], atol=3e-2, rtol=3e-2) and torch.equal(rank, rankings_per_attn[i])):
        raise AssertionError()

from transformers import AutoConfig, AutoModel, AutoTokenizer


tokenizer = AutoTokenizer.from_pretrained("google/embeddinggemma-300m")
config = AutoConfig.from_pretrained("google/embeddinggemma-300m")
config.sliding_window = 6
model = AutoModel.from_pretrained("google/embeddinggemma-300m", config=config)

query = "Which planet is known as the Red Planet?" * 2
inputs = tokenizer(query, return_tensors="pt")

model(**inputs)

from transformers.masking_utils import tensor_to_mask_visual

print(tensor_to_mask_visual(causal_mask_mapping["sliding_attention"][0][0], grid_size=(30, 50)))