File size: 2,406 Bytes
8201ebf f0e7fe0 8201ebf | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | 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",
)
|