import os import accelerate import einops import huggingface_hub import torch import transformers from jinja2 import Environment, FileSystemLoader from llm_studio.app_utils.sections.chat import load_cfg_model_tokenizer from llm_studio.app_utils.utils import hf_repo_friendly_name, save_hf_yaml, set_env from llm_studio.src.utils.config_utils import NON_GENERATION_PROBLEM_TYPES from llm_studio.src.utils.modeling_utils import check_disk_space def get_model_card(cfg, model, repo_id) -> huggingface_hub.ModelCard: """ Method to define the Model Card. It is possible to change the language, the library name, and the tags. These values will appear in the Model Card tab of Hugging Face. Parameters: cfg : Configuration parameters for the model card. model : The model for which the model card is being generated. repo_id : The ID of the target Hugging Face repository. Returns: huggingface_hub.ModelCard: The Model Card containing model information. """ card_data = huggingface_hub.ModelCardData( language="en", library_name="transformers", tags=["gpt", "llm", "large language model", "h2o-llmstudio"], ) cfg_kwargs = dict( use_fast=cfg.tokenizer.use_fast, text_prompt_start=cfg.dataset.text_prompt_start, text_answer_separator=cfg.dataset.text_answer_separator, trust_remote_code=cfg.environment.trust_remote_code, end_of_sentence=( cfg._tokenizer_eos_token if cfg.dataset.add_eos_token_to_prompt else "" ), ) if cfg.problem_type not in NON_GENERATION_PROBLEM_TYPES: cfg_kwargs.update( dict( min_new_tokens=cfg.prediction.min_length_inference, max_new_tokens=cfg.prediction.max_length_inference, do_sample=cfg.prediction.do_sample, num_beams=cfg.prediction.num_beams, temperature=cfg.prediction.temperature, repetition_penalty=cfg.prediction.repetition_penalty, ) ) card = huggingface_hub.ModelCard.from_template( card_data, template_path=os.path.join("model_cards", cfg.environment._model_card_template), base_model=cfg.llm_backbone, # will be replaced in template if it exists repo_id=repo_id, model_architecture=model.backbone.__repr__(), config=cfg.__repr__(), transformers_version=transformers.__version__, einops_version=einops.__version__, accelerate_version=accelerate.__version__, torch_version=torch.__version__.split("+")[0], **cfg_kwargs, ) return card def publish_model_to_hugging_face( path_to_experiment: str, model_name: str, user_id: str = None, api_key: str = None, device: str = "cuda:0", safe_serialization: bool = True, ) -> None: """ Method to publish the model to Hugging Face. Parameters: path_to_experiment: The file path of the fine-tuned model's files. device: The target device for running the model, either 'cpu', 'cpu_shard' or 'cuda:0'. user_id: The Hugging Face user ID. api_key: The Hugging Face API Key. model_name: The name of the model to be published on Hugging Face. safe_serialization: A flag indicating whether safe serialization should be used. Returns: None. The model is published to the specified Hugging Face repository. """ # Check if the 'device' value is valid, raise an exception if not if device == "cpu" or device == "cpu_shard": pass # 'cpu' is a valid value elif device.startswith("cuda:") and device[5:].isdigit(): pass # 'cuda:integer' format is valid else: raise ValueError( "Invalid device value. Use 'cpu', 'cpu_shard' or 'cuda:INTEGER'." ) with set_env(HUGGINGFACE_TOKEN=api_key): cfg, model, tokenizer = load_cfg_model_tokenizer( path_to_experiment, merge=True, device=device, ) check_disk_space(model.backbone, "./") # Check if the user is already logged in, and if not, prompt for API key if api_key: huggingface_hub.login(api_key) # If 'user_id' argument is blank, fetch 'user_id' from the logged-in user if user_id == "": user_id = huggingface_hub.whoami()["name"] repo_id = f"{user_id}/{hf_repo_friendly_name(model_name)}" # push tokenizer to hub tokenizer.push_to_hub(repo_id=repo_id, private=True) # push model card to hub card = get_model_card(cfg, model, repo_id) card.push_to_hub( repo_id=repo_id, repo_type="model", commit_message="Upload model card" ) api = huggingface_hub.HfApi() # push classification head to hub if os.path.isfile(f"{path_to_experiment}/classification_head.pth"): api.upload_file( path_or_fileobj=f"{path_to_experiment}/classification_head.pth", path_in_repo="classification_head.pth", repo_id=repo_id, repo_type="model", commit_message="Upload classification_head.pth", ) # push config to hub api.upload_file( path_or_fileobj=os.path.join(path_to_experiment, "cfg.yaml"), path_in_repo="cfg.yaml", repo_id=repo_id, repo_type="model", commit_message="Upload cfg.yaml", ) # push model to hub model.backbone.config.custom_pipelines = { "text-generation": { "impl": "h2oai_pipeline.H2OTextGenerationPipeline", "pt": "AutoModelForCausalLM", } } model.backbone.push_to_hub( repo_id=repo_id, private=True, commit_message="Upload model", safe_serialization=safe_serialization, ) # Storing HF attributes output_directory = cfg.output_directory save_hf_yaml( path=f"{output_directory.rstrip('/')}/hf.yaml", account_name=user_id, model_name=model_name, repo_id=repo_id, ) # push pipeline to hub template_env = Environment(loader=FileSystemLoader(searchpath="llm_studio/src/")) pipeline_template = template_env.get_template("h2oai_pipeline_template.py") data = { "text_prompt_start": cfg.dataset.text_prompt_start, "text_answer_separator": cfg.dataset.text_answer_separator, } if cfg.dataset.add_eos_token_to_prompt: data.update({"end_of_sentence": cfg._tokenizer_eos_token}) else: data.update({"end_of_sentence": ""}) custom_pipeline = pipeline_template.render(data) custom_pipeline_path = os.path.join(path_to_experiment, "h2oai_pipeline.py") with open(custom_pipeline_path, "w") as f: f.write(custom_pipeline) api.upload_file( path_or_fileobj=custom_pipeline_path, path_in_repo="h2oai_pipeline.py", repo_id=repo_id, repo_type="model", commit_message="Upload h2oai_pipeline.py", )