gem-search/GEM-4B

Model Summary

GEM is a Generative Embedding Model that unifies generation and embedding within a single model. GEM augments retrieval by explicitly reasoning about user intent and relevance criteria, and it is trained to produce reasoning-aligned embeddings.

• Paper: https://arxiv.org/abs/2608.13200
• GitHub: https://github.com/ZhlSh17/GEM

Usage

0. Load the model

from transformers import AutoModel

gem = AutoModel.from_pretrained(
    "gem-search/GEM-4B",
    device_map="auto",
    trust_remote_code=True,  # required
    dtype="bfloat16",
    # attn_implementation="flash_attention_2",  # recommended
).eval()

Make sure to set trust_remote_code=True to use our custom modeling code.

1. Generation

While GEM is an embedding model, it remains generative and can be prompted just like its backbone LLM. For example:

prompt = "Hi, how are you?"
tokenizer = gem.tokenizer
text = tokenizer.apply_chat_template(
    [{"role": "user", "content": prompt}],
    tokenize=False,
    add_generation_prompt=True,
)

inputs = tokenizer([text], return_tensors="pt").to(gem.device)
outputs = gem.generate(**inputs, max_new_tokens=64, do_sample=False)
print(tokenizer.decode(outputs[0]))

2. Embedding

2.1 Generate-then-encode

Use generate_then_encode (or encode_queries) for reasoning-augmented embedding.

query = '''"There is always a day when the flower blooms again."
Find me some similar poems.'''

# if you only need the embedding
q_embs = gem.generate_then_encode([query], max_new_tokens=512, do_sample=False)

# optionally, set `return_generation=True` to return the response
q_embs, resps = gem.generate_then_encode([query],  max_new_tokens=512, do_sample=False, return_generation=True)
print(resps[0])

Note that generate_then_encode accepts additional kwargs for generation, e.g. do_sample, max_new_tokens, and temperature, just like the standard generate.

Alternatively, we recommend using the higher-level API, encode_queries, which calls generate_then_encode internally and supports all of its arguments, plus additional options.

query = '''"There is always a day when the flower blooms again."
Find me some similar poems.'''

q_embs = gem.encode_queries([query], max_new_tokens=512, do_sample=False)

2.2 Embed-only

Use encode_corpus to directly encode texts without generation.

doc = '''No winter stays forever.
Somewhere beneath the frost,
a quiet seed is already
dreaming of spring.'''

doc_embs = gem.encode_corpus([doc])

# all output embeddings are normalized; ready for inner product.
print(q_embs @ doc_embs.T)

2.3 Use custom prompt

In GEM, embedding can be viewed as an additional skill learned by the backbone LLM; the model remains promptable.

By default, GEM uses a reasoning prompt. You may explicitly set your own prompt based on your tasks. Here is an example to use GEM for conventional query expansion:

prompt = """# Task
Write a poem in response to the query below, using no more than 50 words.

# Query
{text}
"""

q_embs, resps = gem.encode_queries([query], prompt=prompt, max_new_tokens=128, do_sample=False, return_generation=True)

print(resps[0])
print(q_embs @ doc_embs.T)

PS: your prompt must contain a {text} placeholder.

Citation

@misc{shen2026gemgenerativeembeddingmodel,
      title={GEM: A Generative Embedding Model Bridging Reasoning and Retrieval}, 
      author={Zhili Shen and Craig Macdonald},
      year={2026},
      eprint={2608.13200},
      archivePrefix={arXiv},
      primaryClass={cs.CL},
      url={https://arxiv.org/abs/2608.13200}, 
}