STM — Gemma-2B

Modular Expert Merging for Biomedical Retrieval — Paper · Collection

Biomedical dense retriever based on google/gemma-2b. Four domain-specific experts were merged with Task Arithmetic.

How to use

# pip install torch==2.6.0 transformers==4.57.1 sentence-transformers==4.1.0 flash-attn --no-build-isolation

import torch
from sentence_transformers import SentenceTransformer
from sentence_transformers.models import Transformer, Pooling

repo_id = "ikim-uk-essen/stm_gemma"

word_embedding = Transformer(
    model_name_or_path=repo_id,
    max_seq_length=512,
    tokenizer_args={"add_eos_token": True},
    model_args=dict(
        dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
    ),
)
pooling = Pooling(2048, pooling_mode="lasttoken")
model = SentenceTransformer(modules=[word_embedding, pooling])

if model.tokenizer.pad_token is None:
    model.tokenizer.pad_token = model.tokenizer.eos_token
    model.tokenizer.padding_side = "left"

task = "Given a question, retrieve relevant passages that answer the question"
query   = f"{task}\nQuery: What are the side effects of metformin?"
passage = "Represent this passage\npassage: Metformin can cause lactic acidosis in rare cases."

emb = model.encode([query, passage], normalize_embeddings=False)
score = float(emb[0] @ emb[1])

Citation

@misc{khattab2026modularexpertmergingbiomedical,
  title         = {Modular Expert Merging for Biomedical Retrieval},
  author        = {Sameh Khattab and Jean-Philippe Corbeil and Osman Alperen {\c{C}}inar-Kora{\c{s}} and Amin Dada and Julian Friedrich and Jiawei He and Douglas Teodoro and Jens Kleesiek},
  year          = {2026},
  eprint        = {2602.04731},
  archivePrefix = {arXiv},
  primaryClass  = {cs.CL},
  url           = {https://arxiv.org/abs/2602.04731}
}