Update convert.py

convert.py

@@ -13,7 +13,6 @@ 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
 
-# Descripción para el reporte en formato Markdown
 REPORT_DESCRIPTION = """
 Este es un reporte automatizado creado con una herramienta de conversión personalizada.
 
@@ -25,52 +24,35 @@ https://colab.research.google.com/github/huggingface/notebooks/blob/main/safeten
 
 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/Hhblvjgvg/convert/discussions
+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.
 """
 
-# Tipo de resultado de conversión: Lista de archivos convertidos y lista de errores
 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]]:
-    """
-    Elimina nombres duplicados en el state_dict basándose en las preferencias y nombres a descartar.
-    """
+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:
-                # Forzar contigüidad
                 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."
-                )
-
+                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]
-
-        # Mecanismo para seleccionar preferentemente claves a mantener provenientes del archivo en disco
         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:
@@ -81,66 +63,34 @@ def _remove_duplicate_names(
     return to_remove
 
 def get_discard_names(model_id: str, revision: Optional[str], folder: str, token: Optional[str]) -> List[str]:
-    """
-    Obtiene los nombres de pesos atados que deben ser descartados según la arquitectura del modelo.
-    """
     try:
         import transformers
-
-        config_filename = hf_hub_download(
-            model_id, revision=revision, filename="config.json", token=token, cache_dir=folder
-        )
+        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)
-
-        # Nombre para esta variable depende de la versión de transformers.
         discard_names = getattr(class_, "_tied_weights_keys", [])
-
     except Exception:
         discard_names = []
     return discard_names
 
 def check_file_size(sf_filename: str, pt_filename: str):
-    """
-    Verifica que la diferencia de tamaño entre el archivo safetensors y el original sea menor al 1%.
-    """
     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%:
-         - {sf_filename}: {sf_size} bytes
-         - {pt_filename}: {pt_size} bytes
-         """
-        )
+        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:
-    """
-    Renombra el archivo PyTorch a safetensors usando el model_id para un mapeo automático.
-    """
     filename, ext = os.path.splitext(pt_filename)
-    # Extraer el nombre base del archivo sin directorios
     base_name = os.path.basename(filename)
-    # Generar el nombre safetensors basado en el model_id y el nombre base
     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:
-    """
-    Convierte modelos con múltiples archivos de pesos (multi-file).
-    """
-    filename = hf_hub_download(
-        repo_id=model_id, revision=revision, filename="pytorch_model.bin.index.json", token=token, cache_dir=folder
-    )
+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 = []
@@ -153,8 +103,6 @@ def convert_multi(
             local_filenames.append(sf_filepath)
         except Exception as e:
             errors.append((filename, e))
-
-    # Crear el archivo de índice para safetensors
     index = os.path.join(folder, f"{model_id.replace('/', '_')}_model.safetensors.index.json")
     try:
         with open(index, "w") as f:
@@ -165,19 +113,11 @@ def convert_multi(
         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:
-    """
-    Convierte un modelo con un único archivo de pesos.
-    """
+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
-        )
+        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)
@@ -188,28 +128,18 @@ def convert_single(
         errors = [("pytorch_model.bin", e)]
     return local_filenames, errors
 
-def convert_file(
-    pt_filename: str,
-    sf_filename: str,
-    discard_names: List[str],
-):
-    """
-    Convierte un archivo de pesos de PyTorch a safetensors.
-    """
+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]
-    # Forzar que los tensores sean contiguos
     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)
@@ -221,27 +151,17 @@ def convert_file(
         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:
-    """
-    Convierte modelos que no utilizan la librería Transformers o que tienen una estructura genérica.
-    """
+def convert_generic(model_id: str, *, revision: Optional[str], folder: str, filenames: Set[str], token: Optional[str]) -> ConversionResult:
     local_filenames = []
     errors = []
-
-    # Agregar ".pth" a las extensiones soportadas
     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
-                )
+                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"}:
-                    # Manejar casos especiales para transformers
                     sf_in_repo = rename(model_id, raw_filename)
                 else:
                     sf_in_repo = rename(model_id, filename)
@@ -252,21 +172,9 @@ def convert_generic(
                 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
-):
-    """
-    Prepara los archivos adicionales necesarios en el repositorio de destino.
-    Descarga o crea archivos como .gitattributes, LICENSE.txt, README.md, USE_POLICY.md, config.json,
-    generation_config.json, special_tokens_map.json, tokenizer.json, tokenizer_config.json.
-    """
+def prepare_target_repo_files(model_id: str, revision: Optional[str], folder: str, token: str, repo_dir: str):
     api = HfApi()
     try:
-        # Descargar archivos comunes del modelo original
         common_files = [
             ".gitattributes",
             "LICENSE.txt",
@@ -280,70 +188,44 @@ def prepare_target_repo_files(
         ]
         for file in common_files:
             try:
-                file_path = hf_hub_download(
-                    repo_id=model_id,
-                    revision=revision,
-                    filename=file,
-                    token=token,
-                    cache_dir=folder
-                )
+                file_path = hf_hub_download(repo_id=model_id, revision=revision, filename=file, token=token, cache_dir=folder)
                 shutil.copy(file_path, repo_dir)
-                print(f"Archivo descargado y copiado: {file}")
-            except Exception as e:
-                # Si el archivo no existe en el modelo original, crear uno vacío o con contenido por defecto
+            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)
-                    print(f"Archivo creado: {file} con configuración para Git LFS")
                 elif file == "LICENSE.txt":
-                    # Crear un archivo LICENSE.txt genérico o personalizado
                     default_license = "MIT License\n\nCopyright (c) 2024"
                     with open(os.path.join(repo_dir, file), "w") as f:
                         f.write(default_license)
-                    print(f"Archivo creado: {file} con licencia por defecto")
                 elif file == "README.md":
-                    # Crear un README.md genérico
                     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)
-                    print(f"Archivo creado: {file} con contenido básico de README")
                 elif file == "USE_POLICY.md":
-                    # Crear un USE_POLICY.md genérico
                     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)
-                    print(f"Archivo creado: {file} con política de uso por defecto")
                 elif file in {"config.json", "generation_config.json", "special_tokens_map.json", "tokenizer.json", "tokenizer_config.json"}:
-                    # Crear archivos JSON vacíos o con contenido por defecto
                     default_json_content = {}
                     with open(os.path.join(repo_dir, file), "w") as f:
                         json.dump(default_json_content, f, indent=4)
-                    print(f"Archivo creado: {file} con contenido JSON vacío")
-                else:
-                    print(f"Error al manejar el archivo {file}: {e}")
     except Exception as e:
-        print(f"Error al preparar archivos adicionales: {e}")
         raise e
 
 def generate_report(model_id: str, local_filenames: List[str], errors: List[Tuple[str, Exception]], output_md_path: str):
-    """
-    Genera un reporte en formato Markdown y JSON detallando los resultados de la conversión.
-    """
-    # Generar reporte Markdown
     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:
@@ -351,14 +233,9 @@ def generate_report(model_id: str, local_filenames: List[str], errors: List[Tupl
             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)
-
-    # Guardar reporte Markdown
     with open(output_md_path, "w") as f:
         f.write(report_content_md)
-
-    # Generar reporte JSON
     report_json = {
         "model_id": model_id,
         "timestamp": datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
@@ -366,31 +243,20 @@ def generate_report(model_id: str, local_filenames: List[str], errors: List[Tupl
         "errors": [{"file": os.path.basename(f), "error": str(e)} for f, e in errors],
         "description": REPORT_DESCRIPTION.strip()
     }
-
-    # Guardar reporte JSON
     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:
-    """
-    Orquesta la conversión del modelo especificado.
-    """
+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 = []
-
-        # Verificar si ya existen archivos .safetensors y si no forzar
         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:
@@ -398,32 +264,22 @@ def convert(
             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
-                    )
+                    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
-                    )
+                    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
-                )
+                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]:
-    """
-    Lee el token de autenticación desde un archivo o variable de entorno.
-    """
     if token_file:
         if os.path.isfile(token_file):
             with open(token_file, "r") as f:
@@ -433,74 +289,40 @@ def read_token(token_file: Optional[str]) -> Optional[str]:
             print(f"El archivo de token especificado no existe: {token_file}")
             return None
     else:
-        # Leer desde variable de entorno
         return os.getenv("HF_TOKEN")
 
 def create_target_repo(model_id: str, api: HfApi, token: str) -> str:
-    """
-    Crea un nuevo repositorio en Hugging Face Hub bajo tu perfil para almacenar los archivos safetensors.
-    """
-    # Definir el nombre del nuevo repositorio
     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
-        )
+        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]):
-    """
-    Sube los archivos convertidos y archivos adicionales a un nuevo repositorio en Hugging Face Hub.
-    """
-    # Inicializar Repository
     repo_dir = "./temp_repo"
     if os.path.exists(repo_dir):
         shutil.rmtree(repo_dir)
     os.makedirs(repo_dir, exist_ok=True)
-
     try:
-        # Clonar el repositorio vacío (crear uno nuevo)
-        repo = Repository(
-            local_dir=repo_dir,
-            clone_from=target_repo_id,
-            use_auth_token=token
-        )
-
-        # Copiar archivos convertidos al repositorio local
+        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)
-
-        # Copiar archivos adicionales al repositorio local
         for file_path in additional_files:
             shutil.copy(file_path, repo_dir)
-
-        # Añadir y hacer commit de los archivos
         repo.git_add(auto_lfs_track=True)
         repo.git_commit("Añadiendo archivos safetensors convertidos")
-        # Push al repositorio
         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:
-        # Limpiar el directorio temporal del repositorio
         shutil.rmtree(repo_dir)
 
 def main():
-    """
-    Función principal que maneja la interacción con el usuario y coordina la conversión, subida y generación de reportes.
-    """
     DESCRIPTION = """
     Herramienta de utilidad simple para convertir automáticamente algunos pesos en el hub al formato `safetensors`.
     Actualmente exclusiva para PyTorch.
@@ -547,14 +369,10 @@ def main():
     )
     args = parser.parse_args()
     model_id = args.model_id
-
-    # Leer el token de autenticación
     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
-
-    # Inicializar HfApi con el token
     api = HfApi()
     try:
         user_info = api.whoami(token=token)
@@ -562,8 +380,6 @@ def main():
     except Exception as e:
         print(f"No se pudo autenticar con Hugging Face Hub: {e}")
         return
-
-    # Confirmación de seguridad
     if args.y:
         proceed = True
     else:
@@ -573,48 +389,27 @@ def main():
             " ¿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)
-
-                # Realizar la conversión
                 local_filenames, errors = convert(model_id, revision=args.revision, force=args.force, token=token)
-
-                # Crear el repositorio de destino en tu perfil
                 target_repo_id = create_target_repo(model_id, api, token)
-
-                # Preparar archivos adicionales en el repositorio local
                 with TemporaryDirectory() as repo_temp_dir:
                     prepare_target_repo_files(model_id, args.revision, folder, token, repo_temp_dir)
-
-                    # Obtener la lista de archivos adicionales
                     additional_files = [os.path.join(repo_temp_dir, f) for f in os.listdir(repo_temp_dir)]
-
-                    # Subir los archivos convertidos y adicionales al repositorio de destino
                     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.")
-
-                # Definir la ruta de salida para el reporte
                 output_md = args.output
                 if args.output_json:
                     output_json = args.output_json
                 else:
                     output_json = os.path.splitext(output_md)[0] + "_report.json"
-
-                # Generar el reporte
                 generate_report(model_id, local_filenames, errors, output_md)
-
-                # Generar reporte JSON adicional si se especificó
-                if args.output_json:
-                    # Ya se ha generado en `generate_report`
-                    pass
-
         except Exception as e:
             print(f"Ocurrió un error inesperado: {e}")
     else: