|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
import argparse |
|
|
import json |
|
|
import os |
|
|
|
|
|
from typing import Any, Dict, Optional |
|
|
|
|
|
import transformers |
|
|
|
|
|
""" |
|
|
Create a randomly initialized HuggingFace model for testing purposes. |
|
|
|
|
|
Model can be specified by name or path for creating its config and tokenizer using |
|
|
HuggingFace transformers AutoConfig and AutoTokenizer functions. |
|
|
|
|
|
Parameter config_updates can be used to override specific model config fields to make |
|
|
it smaller, for example, by changing number of layers or hidden layers dimensionality, |
|
|
making it adequate for testing purposes. This parameter should be specified as |
|
|
a dictionary that can be parsed using json.loads method. |
|
|
|
|
|
Example usage for Llama2 model (requires HF login): |
|
|
``` |
|
|
python tests/setup/models/create_tiny_hf_model.py \ |
|
|
--model_name_or_path meta-llama/Llama-2-7b-hf \ |
|
|
--output_dir tiny_llama2_hf \ |
|
|
--config_updates '{"hidden_size": 128, "num_attention_heads": 4, "num_hidden_layers": 2, "num_key_value_heads": 4}' |
|
|
``` |
|
|
""" |
|
|
|
|
|
|
|
|
def get_hf_model_class(hf_config): |
|
|
"""Get HuggingFace model class from config.""" |
|
|
if len(hf_config.architectures) > 1: |
|
|
print(f"More than one model architecture available, choosing 1st: {hf_config.architectures}") |
|
|
model_name = hf_config.architectures[0] |
|
|
model_class = getattr(transformers, model_name) |
|
|
return model_class |
|
|
|
|
|
|
|
|
def create_hf_model( |
|
|
model_name_or_path: str, output_dir: str, config_updates: Optional[Dict[str, Any]] = None, overwrite: bool = False |
|
|
): |
|
|
"""Create HuggingFace model with optional config updates.""" |
|
|
if os.path.isdir(output_dir) and not overwrite: |
|
|
print(f"Output directory {output_dir} exists and overwrite flag is not set so exiting.") |
|
|
return |
|
|
|
|
|
hf_config = transformers.AutoConfig.from_pretrained(model_name_or_path) |
|
|
tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path) |
|
|
model_class = get_hf_model_class(hf_config) |
|
|
|
|
|
if config_updates is not None: |
|
|
hf_config.update(config_updates) |
|
|
print(hf_config) |
|
|
|
|
|
model = model_class(hf_config) |
|
|
print(model) |
|
|
|
|
|
os.makedirs(output_dir, exist_ok=True) |
|
|
print(f"Saving model to {output_dir}...") |
|
|
tokenizer.save_pretrained(output_dir) |
|
|
model.save_pretrained(output_dir) |
|
|
print("OK.") |
|
|
|
|
|
|
|
|
if __name__ == "__main__": |
|
|
parser = argparse.ArgumentParser("Create a HuggingFace model (random initialization) for testing purposes.") |
|
|
parser.add_argument( |
|
|
"--model_name_or_path", |
|
|
required=True, |
|
|
help="Model name or local path with model config and tokenizer", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--output_dir", |
|
|
required=True, |
|
|
help="Output directory", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--config_updates", |
|
|
type=json.loads, |
|
|
help="Parameter updates in JSON format to overwrite for model config", |
|
|
) |
|
|
parser.add_argument( |
|
|
"--overwrite", |
|
|
action="store_true", |
|
|
help="Overwrite file if it exists", |
|
|
) |
|
|
args = parser.parse_args() |
|
|
create_hf_model(args.model_name_or_path, args.output_dir, args.config_updates, args.overwrite) |
|
|
|
