Matryoshka Representation Learning
Paper • 2205.13147 • Published • 27
How to use mihaylovnikitos/embeddinggemma-300m-hotel with sentence-transformers:
from sentence_transformers import SentenceTransformer
model = SentenceTransformer("mihaylovnikitos/embeddinggemma-300m-hotel")
sentences = [
"task: search result | query: с интернетом",
"title: none | text: бутик-отель в городе: стильные номера, каминная зона, платные услуги хранения багажа на ресепшн, чистая прачечная для гостей. нет специализированных полок, сушилок для ботинок или запираемых ячеек для снаряжения, инфраструктура ориентирована на туристов без спортивного снаряжения.",
"title: none | text: молодёжный хостел 'палитра' в центре города предлагает общие душевые и один открытый бассейн на крыше с общим доступом, кафе с международной кухней, где подают алкоголь и коктейли; формат — общие комнаты, вечерние мероприятия и экскурсии по городу. сегмент — бюджетный.",
"title: none | text: пляжный мини‑отель бюджетного сегмента. бесплатный wi‑fi в лобби и на ресепшн со скоростью до 50 мбит/с, в номерах сигнал нестабилен из‑за толщины стен. номера оборудованы удобными кроватями, завтраки включены, рабочих столов нет."
]
embeddings = model.encode(sentences)
similarities = model.similarity(embeddings, embeddings)
print(similarities.shape)
# [4, 4]This is a sentence-transformers model finetuned from google/embeddinggemma-300m. It maps sentences & paragraphs to a 768-dimensional dense vector space and can be used for retrieval.
SentenceTransformer(
(0): Transformer({'transformer_task': 'feature-extraction', 'modality_config': {'text': {'method': 'forward', 'method_output_name': 'last_hidden_state'}}, 'module_output_name': 'token_embeddings', 'architecture': 'Gemma3TextModel'})
(1): Pooling({'embedding_dimension': 768, 'pooling_mode': 'mean', 'include_prompt': True})
(2): Dense({'in_features': 768, 'out_features': 3072, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'module_input_name': 'sentence_embedding', 'module_output_name': 'sentence_embedding'})
(3): Dense({'in_features': 3072, 'out_features': 768, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'module_input_name': 'sentence_embedding', 'module_output_name': 'sentence_embedding'})
(4): Normalize({})
)
First install the Sentence Transformers library:
pip install -U sentence-transformers
Then you can load this model and run inference.
from sentence_transformers import SentenceTransformer
# Download from the 🤗 Hub
model = SentenceTransformer("mihaylovnikitos/embeddinggemma-300m-hotel")
# Run inference
queries = [
'task: search result | query: быстрый интернет',
]
documents = [
'title: none | text: пляжный хостел бюджетного сегмента: бесплатный wi‑fi в общих кухне и лаунже с ограниченной скоростью до 20–50 мбит/с, симпатичные общие зоны, камеры хранения для сёрферов. рабочих мест и ethernet в номерах нет.',
"title: none | text: речная усадьба 'ледяной берег' — загородный дом на берегу реки с летним шатром для банкетов, 12 уединённых коттеджей, зона барбекю и прокат лодок. подойдет для летних свадеб и выездных церемоний; зимой комплекс закрыт, активностей для катания нет. сегмент — средний.",
"title: none | text: гостевой дом 'старый порт' в центре прибрежного городка предлагает отдельные бюджетные номера и общую кухню, подходит для индивидуальных путешественников и небольших групп. нет приватного двора или пляжной территории; побережье в 10–15 минутах ходьбы, пляжи общественные. услуги ориентированы на краткосрочное проживание и экскурсии по городу. сегмент — бюджетный.",
]
query_embeddings = model.encode_query(queries)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings.shape)
# [1, 768] [3, 768]
# Get the similarity scores for the embeddings
similarities = model.similarity(query_embeddings, document_embeddings)
print(similarities)
# tensor([[0.5166, 0.5017, 0.5486]])
sentence1, sentence2, and label| sentence1 | sentence2 | label | |
|---|---|---|---|
| type | string | string | float |
| modality | text | text | |
| details |
|
|
|
| sentence1 | sentence2 | label |
|---|---|---|
task: search result | query: ищем с девушкой спокойное место у моря, без детских горок и криков, только для взрослых, желательно спа и тихий пляж поблизости — хочется романтики и релакса, на пару дней |
title: none | text: бутик-отель 'морской риф' — adults-only комплекс в 50 метрах от частного мелкогалечного пляжа с зоной для пар и шезлонгами, где не допускаются гости младше 18 лет. на территории современный спа-центр: финская сауна, паровая с соляной стеной, комнаты для парных массажей и процедуры для двоих. номера 'люкс' с панорамными окнами и балконами выходят на море, в каждом — кровать king-size, мини-бар и сервис ужина при свечах по запросу. тихая ухоженная территория, всего 28 номеров, романтические пакеты: ужин на пляже, шампанское и лепестки роз. сегмент — премиум/средний-премиум. |
1.0 |
task: search result | query: ищем с девушкой спокойное место у моря, без детских горок и криков, только для взрослых, желательно спа и тихий пляж поблизости — хочется романтики и релакса, на пару дней |
title: none | text: элегантный курортный отель 'лазурная бухта' расположен на первой линии и предлагает просторный открытый спа-комплекс с бассейном, сауной и расписанием массажей. в отеле есть приватные секции пляжа с зонтиками и баром; в номерах — балконы с видом на море и возможность заказать ужин на террасе. на территории регулярно проходят семейные анимационные программы, работает детский клуб и оборудована игровая площадка; отель принимает гостей любых возрастов. сегмент — средний/семейный. |
0.5 |
task: search result | query: ищем с девушкой спокойное место у моря, без детских горок и криков, только для взрослых, желательно спа и тихий пляж поблизости — хочется романтики и релакса, на пару дней |
title: none | text: бутик-отель 'морской риф' — adults-only комплекс в 50 метрах от частного мелкогалечного пляжа с зоной для пар и шезлонгами, где не допускаются гости младше 18 лет. на территории современный спа-центр: финская сауна, паровая с соляной стеной, комнаты для парных массажей и процедуры для двоих. номера 'люкс' с панорамными окнами и балконами выходят на море, в каждом — кровать king-size, мини-бар и сервис ужина при свечах по запросу. тихая ухоженная территория, всего 28 номеров, романтические пакеты: ужин на пляже, шампанское и лепестки роз. сегмент — премиум/средний-премиум. |
0.7 |
MatryoshkaLoss with these parameters:{
"loss": "CoSENTLoss",
"matryoshka_dims": [
768,
512,
256,
128
],
"matryoshka_weights": [
1,
1,
1,
1
],
"n_dims_per_step": -1
}
sentence1, sentence2, and label| sentence1 | sentence2 | label | |
|---|---|---|---|
| type | string | string | float |
| modality | text | text | |
| details |
|
|
|
| sentence1 | sentence2 | label |
|---|---|---|
task: search result | query: спокойный отдых на море для двоих без детей, adults only, спа, тихий пляж, романтика |
title: none | text: бутик-отель 'морской риф' — adults-only комплекс в 50 метрах от частного мелкогалечного пляжа с зоной для пар и шезлонгами, где не допускаются гости младше 18 лет. на территории современный спа-центр: финская сауна, паровая с соляной стеной, комнаты для парных массажей и процедуры для двоих. номера 'люкс' с панорамными окнами и балконами выходят на море, в каждом — кровать king-size, мини-бар и сервис ужина при свечах по запросу. тихая ухоженная территория, всего 28 номеров, романтические пакеты: ужин на пляже, шампанское и лепестки роз. сегмент — премиум/средний-премиум. |
1.0 |
title: none | text: бутик-отель 'морской риф' — adults-only комплекс в 50 метрах от частного мелкогалечного пляжа с зоной для пар и шезлонгами, где не допускаются гости младше 18 лет. на территории современный спа-центр: финская сауна, паровая с соляной стеной, комнаты для парных массажей и процедуры для двоих. номера 'люкс' с панорамными окнами и балконами выходят на море, в каждом — кровать king-size, мини-бар и сервис ужина при свечах по запросу. тихая ухоженная территория, всего 28 номеров, романтические пакеты: ужин на пляже, шампанское и лепестки роз. сегмент — премиум/средний-премиум. |
title: none | text: пляжный хостел 'морские острова' предлагает недорогие кровати в общих и двухместных номерах, общую кухню и лаунж, бесплатный прокат досок для серфинга и групповые экскурсии. пляж общественный и часто многолюдный; в хостеле молодёжная атмосфера, вечерние игры и совместные выезды. сегмент — бюджетный. |
-1.0 |
task: search result | query: спокойный отдых на море для двоих без детей, adults only, спа, тихий пляж, романтика |
title: none | text: элегантный курортный отель 'лазурная бухта' расположен на первой линии и предлагает просторный открытый спа-комплекс с бассейном, сауной и расписанием массажей. в отеле есть приватные секции пляжа с зонтиками и баром; в номерах — балконы с видом на море и возможность заказать ужин на террасе. на территории регулярно проходят семейные анимационные программы, работает детский клуб и оборудована игровая площадка; отель принимает гостей любых возрастов. сегмент — средний/семейный. |
0.5 |
MatryoshkaLoss with these parameters:{
"loss": "CoSENTLoss",
"matryoshka_dims": [
768,
512,
256,
128
],
"matryoshka_weights": [
1,
1,
1,
1
],
"n_dims_per_step": -1
}
gradient_accumulation_steps: 2num_train_epochs: 1warmup_steps: 0.1data_seed: 42batch_sampler: no_duplicatesdo_predict: Falseprediction_loss_only: Trueper_device_train_batch_size: 8per_device_eval_batch_size: 8gradient_accumulation_steps: 2eval_accumulation_steps: Nonetorch_empty_cache_steps: Nonelearning_rate: 5e-05weight_decay: 0.0adam_beta1: 0.9adam_beta2: 0.999adam_epsilon: 1e-08max_grad_norm: 1.0num_train_epochs: 1max_steps: -1lr_scheduler_type: linearlr_scheduler_kwargs: Nonewarmup_ratio: Nonewarmup_steps: 0.1log_level: passivelog_level_replica: warninglog_on_each_node: Truelogging_nan_inf_filter: Trueenable_jit_checkpoint: Falsesave_on_each_node: Falsesave_only_model: Falserestore_callback_states_from_checkpoint: Falseuse_cpu: Falseseed: 42data_seed: 42bf16: Falsefp16: Falsebf16_full_eval: Falsefp16_full_eval: Falsetf32: Nonelocal_rank: -1ddp_backend: Nonedebug: []dataloader_drop_last: Falsedataloader_num_workers: 0dataloader_prefetch_factor: Nonedisable_tqdm: Falseremove_unused_columns: Truelabel_names: Noneload_best_model_at_end: Falseignore_data_skip: Falsefsdp: []fsdp_config: {'min_num_params': 0, 'xla': False, 'xla_fsdp_v2': False, 'xla_fsdp_grad_ckpt': 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: Nonedeepspeed: Nonelabel_smoothing_factor: 0.0optim: adamw_torch_fusedoptim_args: Nonegroup_by_length: Falselength_column_name: lengthproject: huggingfacetrackio_space_id: trackioddp_find_unused_parameters: Noneddp_bucket_cap_mb: Noneddp_broadcast_buffers: Falsedataloader_pin_memory: Truedataloader_persistent_workers: Falseskip_memory_metrics: Truepush_to_hub: Falseresume_from_checkpoint: Nonehub_model_id: Nonehub_strategy: every_savehub_private_repo: Nonehub_always_push: Falsehub_revision: Nonegradient_checkpointing: Falsegradient_checkpointing_kwargs: Noneinclude_for_metrics: []eval_do_concat_batches: Trueauto_find_batch_size: Falsefull_determinism: Falseddp_timeout: 1800torch_compile: Falsetorch_compile_backend: Nonetorch_compile_mode: Noneinclude_num_input_tokens_seen: noneftune_noise_alpha: Noneoptim_target_modules: Nonebatch_eval_metrics: Falseeval_on_start: Falseuse_liger_kernel: Falseliger_kernel_config: Noneeval_use_gather_object: Falseaverage_tokens_across_devices: Trueuse_cache: Falseprompts: Nonebatch_sampler: no_duplicatesmulti_dataset_batch_sampler: proportionalrouter_mapping: {}learning_rate_mapping: {}| Epoch | Step | Training Loss | Validation Loss |
|---|---|---|---|
| 0.1401 | 25 | 16.1943 | - |
| 0.2801 | 50 | 14.9794 | - |
| 0.4202 | 75 | 14.1160 | - |
| 0.5602 | 100 | 13.5176 | - |
| 0.7003 | 125 | 14.4047 | - |
| 0.8403 | 150 | 13.3981 | - |
| 0.9804 | 175 | 12.8982 | - |
| 1.0 | 179 | - | 10.9653 |
@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",
}
@misc{kusupati2024matryoshka,
title={Matryoshka Representation Learning},
author={Aditya Kusupati and Gantavya Bhatt and Aniket Rege and Matthew Wallingford and Aditya Sinha and Vivek Ramanujan and William Howard-Snyder and Kaifeng Chen and Sham Kakade and Prateek Jain and Ali Farhadi},
year={2024},
eprint={2205.13147},
archivePrefix={arXiv},
primaryClass={cs.LG}
}
@article{10531646,
author={Huang, Xiang and Peng, Hao and Zou, Dongcheng and Liu, Zhiwei and Li, Jianxin and Liu, Kay and Wu, Jia and Su, Jianlin and Yu, Philip S.},
journal={IEEE/ACM Transactions on Audio, Speech, and Language Processing},
title={CoSENT: Consistent Sentence Embedding via Similarity Ranking},
year={2024},
doi={10.1109/TASLP.2024.3402087}
}
Base model
google/embeddinggemma-300m