Embedding / bi-encoder fine-tuning with Unsloth + sentence-transformers by chiggly007 · Pull Request #3718 · unslothai/unsloth

chiggly007 · 2025-12-12T01:27:57Z

Adds docs + runnable example for fine-tuning embedding models with FastModel + LoRA + sentence-transformers. Includes guidance to disable fast-generation for non-causal encoders.

for more information, see https://pre-commit.ci

gemini-code-assist · 2025-12-12T01:28:14Z

Summary of Changes

Hello @chiggly007, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly expands Unsloth's capabilities by enabling and thoroughly documenting the fine-tuning of embedding and bi-encoder models. It provides a comprehensive guide and a runnable example demonstrating how to leverage Unsloth's FastModel and LoRA with the sentence-transformers library. The changes specifically address the unique requirements of non-causal encoder models, such as the necessity to disable fast-generation kernels, allowing users to efficiently fine-tune a broader range of models for tasks like semantic search and information retrieval.

Highlights

Embedding Model Fine-tuning Support: Introduced comprehensive support for fine-tuning embedding and bi-encoder models (e.g., BERT, E5, Arctic-Embed) using Unsloth's FastModel API and LoRA, integrated with the sentence-transformers training ecosystem.
New Documentation Guide: Added a detailed guide (docs/basics/embedding-model-fine-tuning.md) outlining the step-by-step process for embedding model fine-tuning, covering aspects like disabling fast-generation kernels for non-causal encoders, model loading, LoRA adapter configuration, SentenceTransformer wrapping, and training procedures.
Runnable Example Script: Provided a new runnable Python example script (examples/embedding_sentence_transformers.py) that demonstrates an end-to-end workflow for fine-tuning an embedding model, including dataset preparation and training with MultipleNegativesRankingLoss.
README Update: The README.md has been updated with a new section dedicated to embedding/bi-encoder fine-tuning, including direct links to the new example script and the comprehensive guide.
Troubleshooting Guidance: Enhanced the troubleshooting documentation (docs/basics/troubleshooting-and-faqs.md) by adding specific guidance on disabling fast generation for embedding models when encountering CUDA-related runtime errors.
Dataset Shuffling Callback: Implemented a ShuffleDatasetCallback in the example script to ensure reproducible per-epoch randomization of the training dataset, which is crucial for robust model training.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	`/gemini review`	Performs a code review for the current pull request in its current state.
Pull Request Summary	`/gemini summary`	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	`/gemini help`	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

gemini-code-assist

Code Review

This pull request introduces documentation and a runnable example for fine-tuning embedding models using Unsloth and sentence-transformers. The additions are comprehensive and will be very helpful for users. My review focuses on improving the model saving process for correctness and providing a more complete example, along with a minor code style suggestion.

gemini-code-assist · 2025-12-12T01:30:23Z

docs/basics/embedding-model-fine-tuning.md

+sbert_model.save_pretrained("embeddings_merged")
+model.save_pretrained_merged("embeddings_merged", tokenizer)


The current saving method is incorrect as model.save_pretrained_merged will overwrite files created by sbert_model.save_pretrained (like config.json), resulting in a corrupted model directory that cannot be loaded as a SentenceTransformer.

To correctly save the merged SentenceTransformer model, you should first merge the LoRA adapters in-place and then save the sbert_model, which now contains the merged weights. This ensures the entire pipeline (including pooling and normalization layers) is saved correctly.

Suggested change

sbert_model.save_pretrained("embeddings_merged")

model.save_pretrained_merged("embeddings_merged", tokenizer)

model.merge_and_unload()

sbert_model.save_pretrained("embeddings_merged")

gemini-code-assist · 2025-12-12T01:30:23Z

docs/basics/embedding-model-fine-tuning.md

+`sentence-transformers` expects its own `Transformer` module. You can reuse the Unsloth‑loaded model/tokenizer by injecting them into that module.
+
+```python
+import sentence_transformers


The sentence_transformers library is imported locally within the get_st_unsloth_wrapper function. It's better to move all imports to the top of the script for clarity and to follow standard Python conventions. Since sentence_transformers is already imported in the example code block in section 4, this line is redundant and can be removed.

gemini-code-assist · 2025-12-12T01:30:23Z

examples/embedding_sentence_transformers.py

+    tokenizer,
+    base_model_id = BASE_MODEL_ID,
+    pooling_mode = "cls",
+    max_seq_length = MAX_SEQ_LENGTH,


The sentence_transformers library is imported locally within the get_st_unsloth_wrapper function, but it's already imported at the top of the file (lines 24-28). This local import is redundant and should be removed to adhere to Python best practices (PEP 8).

gemini-code-assist · 2025-12-12T01:30:23Z

examples/embedding_sentence_transformers.py

+lora_percentage = round(used_memory_for_lora / max_memory * 100, 3)
+print(
+    f"{round(trainer_stats.metrics['train_runtime'] / 60, 2)} minutes used for training."
+)


model.save_pretrained_merged saves the merged Hugging Face model, but it doesn't save the full SentenceTransformer pipeline (which includes pooling and normalization layers). To provide a more complete and user-friendly example, you can save the entire SentenceTransformer model after merging the LoRA adapters. This allows users to load the complete model with a single command and is more consistent with the goal of fine-tuning a sentence-transformer.

Suggested change

)

print("Merging LoRA adapters...")

model.merge_and_unload()

print(f"Saving merged SentenceTransformer model to {name}_{run}_merged...")

sbert_model.save_pretrained(f"{name}_{run}_merged")

print("Done.")

danielhanchen · 2025-12-12T09:22:50Z

@chiggly007 Oh nice work! Would you be interested in making this inside of Unsloth Docs and a notebook as well?

shimmyshimmer · 2025-12-12T22:45:50Z

Closed because duplicate of: #3719

chiggly007 and others added 2 commits December 11, 2025 20:25

Add embedding/bi-encoder fine-tuning guide and example

41adfb1

[pre-commit.ci] auto fixes from pre-commit.com hooks

efa1be0

for more information, see https://pre-commit.ci

gemini-code-assist bot reviewed Dec 12, 2025

View reviewed changes

chiggly007 added 5 commits December 11, 2025 20:32

Fix embedding model save instructions

2fbbf4a

Remove redundant local sentence-transformers imports

f317e1f

Save merged SentenceTransformer after LoRA merge

80b5636

Add acknowledgements for Arctic-Embed community recipe

30ceb66

Include model card URL in acknowledgements

05016c8

electroglyph mentioned this pull request Dec 12, 2025

add FastSentenceTransformer for easily finetuning SentenceTransformer models #3719

Merged

shimmyshimmer closed this Dec 12, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Embedding / bi-encoder fine-tuning with Unsloth + sentence-transformers#3718

Embedding / bi-encoder fine-tuning with Unsloth + sentence-transformers#3718
chiggly007 wants to merge 7 commits intounslothai:mainfrom
chiggly007:feat/embedding-sentence-transformers

chiggly007 commented Dec 12, 2025

Uh oh!

gemini-code-assist bot commented Dec 12, 2025

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

gemini-code-assist bot Dec 12, 2025

Uh oh!

gemini-code-assist bot Dec 12, 2025

Uh oh!

gemini-code-assist bot Dec 12, 2025

Uh oh!

gemini-code-assist bot Dec 12, 2025

Uh oh!

danielhanchen commented Dec 12, 2025

Uh oh!

shimmyshimmer commented Dec 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		sbert_model.save_pretrained("embeddings_merged")
		model.save_pretrained_merged("embeddings_merged", tokenizer)

-)
+print("Merging LoRA adapters...")
+model.merge_and_unload()
+print(f"Saving merged SentenceTransformer model to {name}_{run}_merged...")
+sbert_model.save_pretrained(f"{name}_{run}_merged")
+print("Done.")

Uh oh!

Conversation

chiggly007 commented Dec 12, 2025

Uh oh!

gemini-code-assist bot commented Dec 12, 2025

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

gemini-code-assist bot Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

gemini-code-assist bot Dec 12, 2025

Choose a reason for hiding this comment

Uh oh!

danielhanchen commented Dec 12, 2025

Uh oh!

shimmyshimmer commented Dec 12, 2025

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants