import argparse import json import os import shutil from collections import defaultdict from datetime import datetime from tempfile import TemporaryDirectory from typing import Dict, List, Optional, Set, Tuple import torch from huggingface_hub import HfApi, Repository, hf_hub_download from huggingface_hub.file_download import repo_folder_name from safetensors.torch import _find_shared_tensors, _is_complete, load_file, save_file REPORT_DESCRIPTION = """ Este es un reporte automatizado creado con una herramienta de conversión personalizada. Este nuevo archivo es equivalente a `pytorch_model.bin` pero es seguro en el sentido de que no se puede inyectar código arbitrario en él. Estos archivos también cargan mucho más rápido que su contraparte de PyTorch: https://colab.research.google.com/github/huggingface/notebooks/blob/main/safetensors_doc/en/speed.ipynb Los widgets en la página de tu modelo funcionarán usando este modelo, asegurando que el archivo realmente funcione. Si encuentras algún problema: por favor repórtalo en el siguiente enlace: https://huggingface.co/spaces/safetensors/convert/discussions Siéntete libre de ignorar este reporte. """ ConversionResult = Tuple[List[str], List[Tuple[str, "Exception"]]] def _remove_duplicate_names(state_dict: Dict[str, torch.Tensor], *, preferred_names: List[str] = None, discard_names: List[str] = None) -> Dict[str, List[str]]: if preferred_names is None: preferred_names = [] preferred_names = set(preferred_names) if discard_names is None: discard_names = [] discard_names = set(discard_names) shareds = _find_shared_tensors(state_dict) to_remove = defaultdict(list) for shared in shareds: complete_names = set([name for name in shared if _is_complete(state_dict[name])]) if not complete_names: if len(shared) == 1: name = list(shared)[0] state_dict[name] = state_dict[name].clone() complete_names = {name} else: raise RuntimeError(f"Error al intentar encontrar nombres para remover al guardar el state dict, pero no se encontró un nombre adecuado para mantener entre: {shared}. Ninguno cubre todo el almacenamiento. Rechazando guardar/cargar el modelo ya que podrías estar almacenando mucha más memoria de la necesaria. Por favor, refiérete a https://huggingface.co/docs/safetensors/torch_shared_tensors para más información. O abre un issue.") keep_name = sorted(list(complete_names))[0] preferred = complete_names.difference(discard_names) if preferred: keep_name = sorted(list(preferred))[0] if preferred_names: preferred = preferred_names.intersection(complete_names) if preferred: keep_name = sorted(list(preferred))[0] for name in sorted(shared): if name != keep_name: to_remove[keep_name].append(name) return to_remove def get_discard_names(model_id: str, revision: Optional[str], folder: str, token: Optional[str]) -> List[str]: try: import transformers config_filename = hf_hub_download(model_id, revision=revision, filename="config.json", token=token, cache_dir=folder) with open(config_filename, "r") as f: config = json.load(f) architecture = config["architectures"][0] class_ = getattr(transformers, architecture) discard_names = getattr(class_, "_tied_weights_keys", []) except Exception: discard_names = [] return discard_names def check_file_size(sf_filename: str, pt_filename: str): sf_size = os.stat(sf_filename).st_size pt_size = os.stat(pt_filename).st_size if (sf_size - pt_size) / pt_size > 0.01: raise RuntimeError(f"La diferencia de tamaño de archivo es mayor al 1%:\n - {sf_filename}: {sf_size} bytes\n - {pt_filename}: {pt_size} bytes") def rename(model_id: str, pt_filename: str) -> str: filename, ext = os.path.splitext(pt_filename) base_name = os.path.basename(filename) safetensors_name = f"{model_id.replace('/', '_')}_{base_name}.safetensors" return safetensors_name def convert_multi(model_id: str, *, revision: Optional[str], folder: str, token: Optional[str], discard_names: List[str]) -> ConversionResult: filename = hf_hub_download(repo_id=model_id, revision=revision, filename="pytorch_model.bin.index.json", token=token, cache_dir=folder) with open(filename, "r") as f: data = json.load(f) filenames = set(data["weight_map"].values()) local_filenames = [] errors = [] for filename in filenames: try: pt_filename = hf_hub_download(repo_id=model_id, filename=filename, token=token, cache_dir=folder) sf_filename = rename(model_id, filename) sf_filepath = os.path.join(folder, sf_filename) convert_file(pt_filename, sf_filepath, discard_names=discard_names) local_filenames.append(sf_filepath) except Exception as e: errors.append((filename, e)) index = os.path.join(folder, f"{model_id.replace('/', '_')}_model.safetensors.index.json") try: with open(index, "w") as f: newdata = {k: v for k, v in data.items()} newmap = {k: rename(model_id, v) for k, v in data["weight_map"].items()} newdata["weight_map"] = newmap json.dump(newdata, f, indent=4) local_filenames.append(index) except Exception as e: errors.append((index, e)) return local_filenames, errors def convert_single(model_id: str, *, revision: Optional[str], folder: str, token: Optional[str], discard_names: List[str]) -> ConversionResult: try: pt_filename = hf_hub_download(repo_id=model_id, revision=revision, filename="pytorch_model.bin", token=token, cache_dir=folder) sf_name = rename(model_id, "pytorch_model.bin") sf_filepath = os.path.join(folder, sf_name) convert_file(pt_filename, sf_filepath, discard_names) local_filenames = [sf_filepath] errors = [] except Exception as e: local_filenames = [] errors = [("pytorch_model.bin", e)] return local_filenames, errors def convert_file(pt_filename: str, sf_filename: str, discard_names: List[str]): loaded = torch.load(pt_filename, map_location="cpu", weights_only=True) if "state_dict" in loaded: loaded = loaded["state_dict"] to_removes = _remove_duplicate_names(loaded, discard_names=discard_names) metadata = {"format": "pt"} for kept_name, to_remove_group in to_removes.items(): for to_remove in to_remove_group: if to_remove not in metadata: metadata[to_remove] = kept_name del loaded[to_remove] loaded = {k: v.contiguous() for k, v in loaded.items()} dirname = os.path.dirname(sf_filename) os.makedirs(dirname, exist_ok=True) save_file(loaded, sf_filename, metadata=metadata) check_file_size(sf_filename, pt_filename) reloaded = load_file(sf_filename) for k in loaded: pt_tensor = loaded[k] sf_tensor = reloaded[k] if not torch.equal(pt_tensor, sf_tensor): raise RuntimeError(f"Los tensores de salida no coinciden para la clave {k}") def convert_generic(model_id: str, *, revision: Optional[str], folder: str, filenames: Set[str], token: Optional[str]) -> ConversionResult: local_filenames = [] errors = [] extensions = set([".bin", ".ckpt", ".pth"]) for filename in filenames: prefix, ext = os.path.splitext(filename) if ext in extensions: try: pt_filename = hf_hub_download(model_id, revision=revision, filename=filename, token=token, cache_dir=folder) dirname, raw_filename = os.path.split(filename) if raw_filename in {"pytorch_model.bin", "pytorch_model.pth"}: sf_in_repo = rename(model_id, raw_filename) else: sf_in_repo = rename(model_id, filename) sf_filepath = os.path.join(folder, sf_in_repo) convert_file(pt_filename, sf_filepath, discard_names=[]) local_filenames.append(sf_filepath) except Exception as e: errors.append((filename, e)) return local_filenames, errors def prepare_target_repo_files(model_id: str, revision: Optional[str], folder: str, token: str, repo_dir: str): api = HfApi() try: common_files = [ ".gitattributes", "LICENSE.txt", "README.md", "USE_POLICY.md", "config.json", "generation_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json" ] for file in common_files: try: file_path = hf_hub_download(repo_id=model_id, revision=revision, filename=file, token=token, cache_dir=folder) shutil.copy(file_path, repo_dir) except Exception: if file == ".gitattributes": gitattributes_content = "model.safetensors filter=safetensors diff=safetensors merge=safetensors -text\n" with open(os.path.join(repo_dir, file), "w") as f: f.write(gitattributes_content) elif file == "LICENSE.txt": default_license = "MIT License\n\nCopyright (c) 2024" with open(os.path.join(repo_dir, file), "w") as f: f.write(default_license) elif file == "README.md": readme_content = f"# {model_id.replace('/', ' ').title()}\n\nModelo convertido a safetensors." with open(os.path.join(repo_dir, file), "w") as f: f.write(readme_content) elif file == "USE_POLICY.md": use_policy_content = "### Política de Uso\n\nEste modelo se distribuye bajo términos de uso estándar." with open(os.path.join(repo_dir, file), "w") as f: f.write(use_policy_content) elif file in {"config.json", "generation_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"}: default_json_content = {} with open(os.path.join(repo_dir, file), "w") as f: json.dump(default_json_content, f, indent=4) except Exception as e: raise e def generate_report(model_id: str, local_filenames: List[str], errors: List[Tuple[str, Exception]], output_md_path: str): report_lines = [ f"# Reporte de Conversión para el Modelo `{model_id}`", f"Fecha y Hora: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}", "", "## Archivos Convertidos Exitosamente", ] if local_filenames: for filename in local_filenames: report_lines.append(f"- `{os.path.basename(filename)}`") else: report_lines.append("No se convirtieron archivos.") report_lines.append("") report_lines.append("## Errores Durante la Conversión") if errors: for filename, error in errors: report_lines.append(f"- **Archivo**: `{os.path.basename(filename)}`\n - **Error**: {error}") else: report_lines.append("No hubo errores durante la conversión.") report_content_md = "\n".join(report_lines) with open(output_md_path, "w") as f: f.write(report_content_md) report_json = { "model_id": model_id, "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S'), "converted_files": [os.path.basename(f) for f in local_filenames], "errors": [{"file": os.path.basename(f), "error": str(e)} for f, e in errors], "description": REPORT_DESCRIPTION.strip() } json_output_path = os.path.splitext(output_md_path)[0] + "_report.json" with open(json_output_path, "w") as f: json.dump(report_json, f, indent=4) print(f"Reportes generados en: {output_md_path} y {json_output_path}") def convert(model_id: str, revision: Optional[str] = None, force: bool = False, token: Optional[str] = None) -> ConversionResult: api = HfApi() info = api.model_info(repo_id=model_id, revision=revision) filenames = set(s.rfilename for s in info.siblings) with TemporaryDirectory() as d: folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) os.makedirs(folder, exist_ok=True) local_filenames = [] errors = [] if not force and any(filename.endswith(".safetensors") for filename in filenames): print(f"El modelo `{model_id}` ya tiene archivos `.safetensors` convertidos. Usando report existente o forzando con --force.") else: library_name = getattr(info, "library_name", None) if library_name == "transformers": discard_names = get_discard_names(model_id, revision=revision, folder=folder, token=token) if "pytorch_model.bin" in filenames or "pytorch_model.pth" in filenames: converted, conv_errors = convert_single(model_id, revision=revision, folder=folder, token=token, discard_names=discard_names) local_filenames.extend(converted) errors.extend(conv_errors) elif "pytorch_model.bin.index.json" in filenames: converted, conv_errors = convert_multi(model_id, revision=revision, folder=folder, token=token, discard_names=discard_names) local_filenames.extend(converted) errors.extend(conv_errors) else: print(f"El modelo `{model_id}` no parece ser un modelo válido de PyTorch. No se puede convertir.") else: converted, conv_errors = convert_generic(model_id, revision=revision, folder=folder, filenames=filenames, token=token) local_filenames.extend(converted) errors.extend(conv_errors) return local_filenames, errors def read_token(token_file: Optional[str]) -> Optional[str]: if token_file: if os.path.isfile(token_file): with open(token_file, "r") as f: token = f.read().strip() return token else: print(f"El archivo de token especificado no existe: {token_file}") return None else: return os.getenv("HF_TOKEN") def create_target_repo(model_id: str, api: HfApi, token: str) -> str: target_repo_id = f"{api.whoami(token=token)['name']}/{model_id.replace('/', '_')}_safetensors" try: api.create_repo(name=f"{model_id.replace('/', '_')}_safetensors", repo_type="model", exist_ok=True, token=token) print(f"Repositorio creado o ya existente: {target_repo_id}") except Exception as e: print(f"Error al crear el repositorio `{target_repo_id}`: {e}") raise e return target_repo_id def upload_to_hf(local_filenames: List[str], target_repo_id: str, token: str, additional_files: List[str]): repo_dir = "./temp_repo" if os.path.exists(repo_dir): shutil.rmtree(repo_dir) os.makedirs(repo_dir, exist_ok=True) try: repo = Repository(local_dir=repo_dir, clone_from=target_repo_id, use_auth_token=token) for file_path in local_filenames: shutil.copy(file_path, repo_dir) for file_path in additional_files: shutil.copy(file_path, repo_dir) repo.git_add(auto_lfs_track=True) repo.git_commit("Añadiendo archivos safetensors convertidos") repo.git_push() print(f"Archivos subidos exitosamente al repositorio: {target_repo_id}") except Exception as e: print(f"Error al subir archivos al repositorio `{target_repo_id}`: {e}") raise e finally: shutil.rmtree(repo_dir) def main(): DESCRIPTION = """ Herramienta de utilidad simple para convertir automáticamente algunos pesos en el hub al formato `safetensors`. Actualmente exclusiva para PyTorch. Funciona descargando los pesos (PT), convirtiéndolos localmente, subiéndolos a tu propio perfil en Hugging Face Hub y generando reportes en formato Markdown y JSON. """ parser = argparse.ArgumentParser(description=DESCRIPTION) parser.add_argument( "model_id", type=str, help="El nombre del modelo en el hub para convertir. Por ejemplo, `gpt2` o `facebook/wav2vec2-base-960h`", ) parser.add_argument( "--revision", type=str, help="La revisión a convertir", ) parser.add_argument( "--force", action="store_true", help="Forzar la conversión incluso si ya existen archivos `.safetensors` en el modelo.", ) parser.add_argument( "-y", action="store_true", help="Ignorar prompt de seguridad", ) parser.add_argument( "--output", type=str, default="conversion_report.md", help="Ruta donde se guardará el reporte de conversión en formato Markdown.", ) parser.add_argument( "--output-json", type=str, default=None, help="Ruta donde se guardará el reporte de conversión en formato JSON. Si no se especifica, se creará en la misma ubicación que el reporte Markdown.", ) parser.add_argument( "--token-file", type=str, default=None, help="Ruta al archivo que contiene el token de autenticación de Hugging Face. Si no se especifica, se intentará leer desde la variable de entorno 'HF_TOKEN'.", ) args = parser.parse_args() model_id = args.model_id token = read_token(args.token_file) if not token: print("No se proporcionó un token de autenticación válido. Por favor, proporciónalo mediante --token-file o establece la variable de entorno 'HF_TOKEN'.") return api = HfApi() try: user_info = api.whoami(token=token) print(f"Autenticado como: {user_info['name']}") except Exception as e: print(f"No se pudo autenticar con Hugging Face Hub: {e}") return if args.y: proceed = True else: txt = input( "Este script de conversión desenpaca un archivo pickled, lo cual es inherentemente inseguro. Si no confías en este archivo, te invitamos a usar " "https://huggingface.co/spaces/safetensors/convert o Google Colab u otra solución alojada para evitar posibles problemas con este archivo." " ¿Continuar [Y/n] ? " ) proceed = txt.lower() in {"", "y", "yes"} if proceed: try: with TemporaryDirectory() as d: folder = os.path.join(d, repo_folder_name(repo_id=model_id, repo_type="models")) os.makedirs(folder, exist_ok=True) local_filenames, errors = convert(model_id, revision=args.revision, force=args.force, token=token) target_repo_id = create_target_repo(model_id, api, token) with TemporaryDirectory() as repo_temp_dir: prepare_target_repo_files(model_id, args.revision, folder, token, repo_temp_dir) additional_files = [os.path.join(repo_temp_dir, f) for f in os.listdir(repo_temp_dir)] if local_filenames or additional_files: upload_to_hf(local_filenames, target_repo_id, token, additional_files) print(f"Archivos convertidos y adicionales subidos exitosamente a: {target_repo_id}") else: print("No hay archivos convertidos ni adicionales para subir.") output_md = args.output if args.output_json: output_json = args.output_json else: output_json = os.path.splitext(output_md)[0] + "_report.json" generate_report(model_id, local_filenames, errors, output_md) except Exception as e: print(f"Ocurrió un error inesperado: {e}") else: print(f"La respuesta fue '{txt}', abortando.") if __name__ == "__main__": main()