mirror of
https://github.com/ml-explore/mlx-examples.git
synced 2025-06-24 09:21:18 +08:00
91 lines
2.6 KiB
Python
91 lines
2.6 KiB
Python
|
# Copyright © 2023 Apple Inc.
|
||
|
|
|
||
|
|
import glob
|
||
|
|
import json
|
||
|
|
from pathlib import Path
|
||
|
|
|
||
|
|
import mlx.core as mx
|
||
|
|
import transformers
|
||
|
|
from huggingface_hub import snapshot_download
|
||
|
|
|
||
|
|
|
||
|
|
def fetch_from_hub(hf_path: str):
|
||
|
|
model_path = snapshot_download(
|
||
|
|
repo_id=hf_path,
|
||
|
|
allow_patterns=["*.json", "*.safetensors", "tokenizer.model"],
|
||
|
|
)
|
||
|
|
weight_files = glob.glob(f"{model_path}/*.safetensors")
|
||
|
|
if len(weight_files) == 0:
|
||
|
|
raise FileNotFoundError("No safetensors found in {}".format(model_path))
|
||
|
|
|
||
|
|
weights = {}
|
||
|
|
for wf in weight_files:
|
||
|
|
weights.update(mx.load(wf).items())
|
||
|
|
|
||
|
|
config = transformers.AutoConfig.from_pretrained(hf_path)
|
||
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(
|
||
|
|
hf_path,
|
||
|
|
)
|
||
|
|
return weights, config.to_dict(), tokenizer
|
||
|
|
|
||
|
|
|
||
|
|
def upload_to_hub(path: str, name: str, hf_path: str):
|
||
|
|
import os
|
||
|
|
|
||
|
|
from huggingface_hub import HfApi, ModelCard, logging
|
||
|
|
|
||
|
|
repo_id = f"mlx-community/{name}"
|
||
|
|
|
||
|
|
card = ModelCard.load(hf_path)
|
||
|
|
card.data.tags = ["mlx"] if card.data.tags is None else card.data.tags + ["mlx"]
|
||
|
|
card.text = f"""
|
||
|
|
# {name}
|
||
|
|
This model was converted to MLX format from [`{hf_path}`]().
|
||
|
|
Refer to the [original model card](https://huggingface.co/{hf_path}) for more details on the model.
|
||
|
|
## Use with mlx
|
||
|
|
```bash
|
||
|
|
pip install mlx
|
||
|
|
git clone https://github.com/ml-explore/mlx-examples.git
|
||
|
|
cd mlx-examples/llms/hf_llm
|
||
|
|
python generate.py --model {repo_id} --prompt "My name is"
|
||
|
|
```
|
||
|
|
"""
|
||
|
|
card.save(os.path.join(path, "README.md"))
|
||
|
|
|
||
|
|
logging.set_verbosity_info()
|
||
|
|
|
||
|
|
api = HfApi()
|
||
|
|
api.create_repo(repo_id=repo_id, exist_ok=True)
|
||
|
|
api.upload_folder(
|
||
|
|
folder_path=path,
|
||
|
|
repo_id=repo_id,
|
||
|
|
repo_type="model",
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def make_shards(weights: dict, max_file_size_gibibyte: int = 15):
|
||
|
|
max_file_size_bytes = max_file_size_gibibyte << 30
|
||
|
|
shards = []
|
||
|
|
shard, shard_size = {}, 0
|
||
|
|
for k, v in weights.items():
|
||
|
|
estimated_size = v.size * v.dtype.size
|
||
|
|
if shard_size + estimated_size > max_file_size_bytes:
|
||
|
|
shards.append(shard)
|
||
|
|
shard, shard_size = {}, 0
|
||
|
|
shard[k] = v
|
||
|
|
shard_size += estimated_size
|
||
|
|
shards.append(shard)
|
||
|
|
return shards
|
||
|
|
|
||
|
|
|
||
|
|
def save_model(save_dir: str, weights, tokenizer, config):
|
||
|
|
save_dir = Path(save_dir)
|
||
|
|
save_dir.mkdir(parents=True, exist_ok=True)
|
||
|
|
shards = make_shards(weights)
|
||
|
|
for i, shard in enumerate(shards):
|
||
|
|
# TODO use HF file name scheme for simplicity
|
||
|
|
mx.save_safetensors(str(save_dir / f"weights.{i:02d}.safetensors"), shard)
|
||
|
|
tokenizer.save_pretrained(save_dir)
|
||
|
|
with open(save_dir / "config.json", "w") as fid:
|
||
|
|
json.dump(config, fid, indent=4)
|