lfm3 / .meta /create_model.py
yujiepan's picture
Upload folder using huggingface_hub
f0e7fe0 verified
Raw History Blame Contribute Delete
2.41 kB
def create_model(source_model_id: str, source_revision: str, save_folder: str) -> str:
import json
from pathlib import Path
import torch
from huggingface_hub import hf_hub_download
from transformers import AutoConfig, AutoModelForCausalLM, set_seed
folder = Path(save_folder)
folder.mkdir(parents=True, exist_ok=True)
# LFM3 is not registered in the installed Transformers release; ship its pinned
# remote code so the checkpoint loads without a newer package.
for name in (
"configuration_lfm3.py",
"modeling_lfm3.py",
"lfm3_conv.py",
"chat_template.jinja",
"generation_config.json",
"special_tokens_map.json",
"tokenizer.json",
"tokenizer_config.json",
"LICENSE",
):
(folder / name).write_bytes(Path(hf_hub_download(source_model_id, name, revision=source_revision)).read_bytes())
config = json.loads(Path(hf_hub_download(source_model_id, "config.json", revision=source_revision)).read_text())
# Keep the two-conv prefix and two complete full/conv/sliding/conv cycles.
# Match tiny-random/lfm2's residual/MLP widths; retain LFM3's 4:1 Q/KV ratio.
config.update(
hidden_size=64,
intermediate_size=128,
num_hidden_layers=10,
num_attention_heads=4,
num_key_value_heads=1,
head_dim=32,
max_position_embeddings=128000,
)
config["layer_types"] = config["layer_types"][:10]
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
set_seed(42)
model = AutoModelForCausalLM.from_config(
AutoConfig.from_pretrained(folder, trust_remote_code=True), trust_remote_code=True, dtype=torch.bfloat16
)
model.save_pretrained(folder)
# `save_pretrained()` rewrites these two files with this environment's
# versions; restore the pinned source bytes plus the deliberate config edits
# so only the intended fields differ from the source.
(folder / "config.json").write_text(json.dumps(config, indent=2) + "\n")
(folder / "generation_config.json").write_bytes(
Path(hf_hub_download(source_model_id, "generation_config.json", revision=source_revision)).read_bytes()
)
return str(model)
create_model(
source_model_id="LiquidAI/LFM3-1.7B",
source_revision="ac2586a34f1a1c03329dbdcdefabd8b6478ddaa7",
save_folder="/tmp/lfm3",
)