ceos/unsloth
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Add push_to_hub_gguf support for FastSentenceTransformer (#4002)

Browse source 
* Implement GGUF upload method for SentenceTransformer

Added a method to convert and upload SentenceTransformer models to GGUF format, including handling of tokenizer, quantization methods, and repository management on Hugging Face Hub.

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

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

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
This commit is contained in:
Etherll 2026-02-09 11:51:26 +03:00 • committed by GitHub
commit 858537610a
1 changed files with 235 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

235
unsloth/models/sentence_transformer.py
View file

@ -288,6 +288,237 @@ def _save_pretrained_gguf(
return result
def _push_to_hub_gguf(
self,
repo_id,
tokenizer = None,
quantization_method = "fast_quantized",
first_conversion = None,
token = None,
private = None,
commit_message = "Upload GGUF SentenceTransformer model trained with Unsloth",
commit_description = "Upload GGUF model trained with Unsloth 2x faster",
max_shard_size = "5GB",
temporary_location = "_unsloth_temporary_saved_buffers",
maximum_memory_usage = 0.85,
create_pr = False,
revision = None,
tags = None,
**kwargs,
):
"""
Converts the SentenceTransformer model to GGUF format and pushes to the Hugging Face Hub.
This method:
1. Saves the model locally to a temporary directory in GGUF format.
2. Uploads the GGUF files, config, Ollama Modelfile, and README to the Hub.
3. Cleans up the temporary directory.
Args:
repo_id (str): The Hugging Face Hub repo ID (e.g., "username/model-name").
tokenizer: The tokenizer to save. Defaults to `self.tokenizer`.
quantization_method (str or list): GGUF quantization method(s). Can be a string or list of strings.
Choose from the following options:
* "not_quantized" : Recommended. Fast conversion. Slow inference, big files.
* "fast_quantized" : Recommended. Fast conversion. OK inference, OK file size.
* "quantized" : Recommended. Slow conversion. Fast inference, small files.
* "f32" : Not recommended. Retains 100% accuracy, but super slow and memory hungry.
* "f16" : Fastest conversion + retains 100% accuracy. Slow and memory hungry.
* "q8_0" : Fast conversion. High resource use, but generally acceptable.
* "q4_k_m" : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q4_K
* "q5_k_m" : Recommended. Uses Q6_K for half of the attention.wv and feed_forward.w2 tensors, else Q5_K
* "q2_k" : Uses Q4_K for the attention.vw and feed_forward.w2 tensors, Q2_K for the other tensors.
* "q3_k_l" : Uses Q5_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
* "q3_k_m" : Uses Q4_K for the attention.wv, attention.wo, and feed_forward.w2 tensors, else Q3_K
* "q3_k_s" : Uses Q3_K for all tensors
* "q4_0" : Original quant method, 4-bit.
* "q4_1" : Higher accuracy than q4_0 but not as high as q5_0. However has quicker inference than q5 models.
* "q4_k_s" : Uses Q4_K for all tensors
* "q5_0" : Higher accuracy, higher resource usage and slower inference.
* "q5_1" : Even higher accuracy, resource usage and slower inference.
* "q5_k_s" : Uses Q5_K for all tensors
* "q6_k" : Uses Q8_K for all tensors
first_conversion (str, optional): The initial conversion format before quantization.
token (str, optional): Hugging Face token. Uses cached token if not provided.
private (bool, optional): Whether the repo should be private.
commit_message (str): Commit message for the upload.
commit_description (str): Commit description for the upload.
max_shard_size (str): Maximum shard size for saving.
temporary_location (str): Temp directory for intermediate files.
maximum_memory_usage (float): Max fraction of memory to use.
create_pr (bool): Whether to create a pull request instead of pushing directly.
revision (str, optional): Branch/revision to push to.
tags (list, optional): Additional tags for the repo.
Returns:
str: The full repo ID on Hugging Face Hub.
"""
if token is None:
token = get_token()
if token is None:
raise ValueError(
"No HF token provided. Please provide a token or login with `huggingface-cli login`"
)
api = HfApi(token = token)
# Determine full repo_id
if "/" not in repo_id:
username = api.whoami()["name"]
full_repo_id = f"{username}/{repo_id}"
else:
full_repo_id = repo_id
model_name = full_repo_id.split("/")[-1]
# Create repo
try:
api.create_repo(
repo_id = full_repo_id,
private = private,
exist_ok = True,
repo_type = "model",
)
except Exception as e:
print(f"Unsloth Warning: Could not create repo: {e}")
# Save to temporary directory first
with tempfile.TemporaryDirectory(prefix = "unsloth_st_gguf_") as temp_dir:
print(f"Unsloth: Converting SentenceTransformer to GGUF format...")
# Call save_pretrained_gguf to do the local conversion
result = _save_pretrained_gguf(
self,
save_directory = temp_dir,
tokenizer = tokenizer,
quantization_method = quantization_method,
first_conversion = first_conversion,
push_to_hub = False, # We handle upload ourselves
token = token,
max_shard_size = max_shard_size,
temporary_location = temporary_location,
maximum_memory_usage = maximum_memory_usage,
)
gguf_files = result.get("gguf_files", [])
modelfile_location = result.get("modelfile_location", None)
is_vlm = result.get("is_vlm", False)
fix_bos_token = result.get("fix_bos_token", False)
print(f"Unsloth: Uploading GGUF to https://huggingface.co/{full_repo_id}...")
# Upload GGUF files
for file_location in gguf_files:
if os.path.exists(file_location):
filename = os.path.basename(file_location)
print(f" Uploading {filename}...")
api.upload_file(
path_or_fileobj = file_location,
path_in_repo = filename,
repo_id = full_repo_id,
repo_type = "model",
commit_message = commit_message,
commit_description = commit_description,
create_pr = create_pr,
revision = revision,
)
# Upload Modelfile if exists
if modelfile_location and os.path.exists(modelfile_location):
print(" Uploading Ollama Modelfile...")
api.upload_file(
path_or_fileobj = modelfile_location,
path_in_repo = "Modelfile",
repo_id = full_repo_id,
repo_type = "model",
commit_message = f"{commit_message} - Ollama Modelfile",
create_pr = create_pr,
revision = revision,
)
# Upload config.json if exists
config_path = os.path.join(temp_dir, "config.json")
if os.path.exists(config_path):
print(" Uploading config.json...")
api.upload_file(
path_or_fileobj = config_path,
path_in_repo = "config.json",
repo_id = full_repo_id,
repo_type = "model",
commit_message = f"{commit_message} - config",
create_pr = create_pr,
revision = revision,
)
# Create and upload README
gguf_basenames = [os.path.basename(f) for f in gguf_files if os.path.exists(f)]
readme_content = f"""---
tags:
- gguf
- llama.cpp
- unsloth
- sentence-transformers
{"- vision-language-model" if is_vlm else ""}
---
# {model_name} - GGUF
This sentence-transformers model was finetuned and converted to GGUF format using [Unsloth](https://github.com/unslothai/unsloth).
## Available Model files:
"""
for fname in gguf_basenames:
readme_content += f"- `{fname}`\n"
if modelfile_location and os.path.exists(modelfile_location):
readme_content += "\n## Ollama\n"
readme_content += "An Ollama Modelfile is included for easy deployment.\n"
if fix_bos_token:
readme_content += "\n## Note\n"
readme_content += (
"The model's BOS token behavior was adjusted for GGUF compatibility.\n"
)
readme_content += (
"\nThis was trained 2x faster with [Unsloth](https://github.com/unslothai/unsloth)\n"
'[<img src="https://raw.githubusercontent.com/unslothai/unsloth/main/images/unsloth%20made%20with%20love.png" width="200"/>](https://github.com/unslothai/unsloth)\n'
)
readme_path = os.path.join(temp_dir, "README.md")
with open(readme_path, "w", encoding = "utf-8") as f:
f.write(readme_content)
api.upload_file(
path_or_fileobj = readme_path,
path_in_repo = "README.md",
repo_id = full_repo_id,
repo_type = "model",
commit_message = "Add README",
create_pr = create_pr,
revision = revision,
)
# Add tags
all_tags = ["gguf", "llama-cpp", "unsloth", "sentence-transformers"]
if is_vlm:
all_tags.append("vision-language-model")
if tags is not None:
if isinstance(tags, (list, tuple)):
all_tags.extend(tags)
else:
all_tags.append(tags)
try:
api.add_tags(repo_id = full_repo_id, tags = all_tags, repo_type = "model")
except:
pass
print(
f"Unsloth: Successfully uploaded GGUF to https://huggingface.co/{full_repo_id}"
)
return full_repo_id
class FastSentenceTransformer(FastModel):
@staticmethod
def _read_pooling_mode(model_name, token):
@ -1321,6 +1552,8 @@ class FastSentenceTransformer(FastModel):
_save_pretrained_gguf, st_model
)
st_model.push_to_hub_gguf = types.MethodType(_push_to_hub_gguf, st_model)
def _push_to_hub_merged(self, repo_id, **push_kwargs):
hub_token = push_kwargs.get("token", None) or get_token()
if hub_token is None:
@ -1523,6 +1756,8 @@ class FastSentenceTransformer(FastModel):
_save_pretrained_gguf, st_model
)
st_model.push_to_hub_gguf = types.MethodType(_push_to_hub_gguf, st_model)
def _push_to_hub_merged(self, repo_id, **kwargs):
token = kwargs.get("token", None) or get_token()
if token is None: