in src/transformers/utils/args_doc.py [0:0]
def auto_class_docstring(cls, custom_intro=None, custom_args=None, checkpoint=None):
"""
Wrapper that automatically generates a docstring for classes based on their attributes and methods.
"""
# import here to avoid circular import
from transformers.models import auto as auto_module
is_dataclass = False
docstring_init = ""
if "PreTrainedModel" in (x.__name__ for x in cls.__mro__):
docstring_init = auto_method_docstring(
cls.__init__, parent_class=cls, custom_args=custom_args
).__doc__.replace("Args:", "Parameters:")
elif "ModelOutput" in (x.__name__ for x in cls.__mro__):
# We have a data class
is_dataclass = True
doc_class = cls.__doc__
if custom_args is None and doc_class:
custom_args = doc_class
docstring_args = auto_method_docstring(
cls.__init__,
parent_class=cls,
custom_args=custom_args,
source_args_dict=get_args_doc_from_source(ModelOutputArgs),
).__doc__
indent_level = get_indent_level(cls)
model_name_lowercase = get_model_name(cls)
model_name_title = " ".join([k.title() for k in model_name_lowercase.split("_")]) if model_name_lowercase else None
if model_name_lowercase and model_name_lowercase not in getattr(
getattr(auto_module, PLACEHOLDER_TO_AUTO_MODULE["config_class"][0]),
PLACEHOLDER_TO_AUTO_MODULE["config_class"][1],
):
model_name_lowercase = model_name_lowercase.replace("_", "-")
name = re.findall(rf"({'|'.join(ClassDocstring.__dict__.keys())})$", cls.__name__)
if name == [] and custom_intro is None and not is_dataclass:
raise ValueError(
f"`{cls.__name__}` is not registered in the auto doc. Here are the available classes: {ClassDocstring.__dict__.keys()}.\n"
"Add a `custom_intro` to the decorator if you want to use `auto_docstring` on a class not registered in the auto doc."
)
if name != [] or custom_intro is not None or is_dataclass:
name = name[0] if name else None
if custom_intro is not None:
pre_block = equalize_indent(custom_intro, indent_level)
if not pre_block.endswith("\n"):
pre_block += "\n"
elif model_name_title is None or name is None:
pre_block = ""
else:
pre_block = getattr(ClassDocstring, name).format(model_name=model_name_title)
# Start building the docstring
docstring = set_min_indent(f"{pre_block}", indent_level) if len(pre_block) else ""
if name != "PreTrainedModel" and "PreTrainedModel" in (x.__name__ for x in cls.__mro__):
docstring += set_min_indent(f"{ClassDocstring.PreTrainedModel}", indent_level)
# Add the __init__ docstring
if docstring_init:
docstring += set_min_indent(f"\n{docstring_init}", indent_level)
elif is_dataclass:
# No init function, we have a data class
docstring += "\nArgs:\n" if not docstring_args else docstring_args
source_args_dict = get_args_doc_from_source(ModelOutputArgs)
doc_class = cls.__doc__ if cls.__doc__ else ""
documented_kwargs, _ = parse_docstring(doc_class)
for param_name, param_type_annotation in cls.__annotations__.items():
param_type = str(param_type_annotation)
optional = False
# Process parameter type
if "typing" in param_type:
param_type = "".join(param_type.split("typing.")).replace("transformers.", "~")
else:
param_type = f"{param_type.replace('transformers.', '~').replace('builtins', '')}.{param_name}"
if "ForwardRef" in param_type:
param_type = re.sub(r"ForwardRef\('([\w.]+)'\)", r"\1", param_type)
if "Optional" in param_type:
param_type = re.sub(r"Optional\[(.*?)\]", r"\1", param_type)
optional = True
# Check for default value
param_default = ""
param_default = str(getattr(cls, param_name, ""))
param_default = f", defaults to `{param_default}`" if param_default != "" else ""
param_type, optional_string, shape_string, additional_info, description, is_documented = (
_get_parameter_info(param_name, documented_kwargs, source_args_dict, param_type, optional)
)
if is_documented:
# Check if type is missing
if param_type == "":
print(f"🚨 {param_name} for {cls.__qualname__} in file {cls.__code__.co_filename} has no type")
param_type = param_type if "`" in param_type else f"`{param_type}`"
# Format the parameter docstring
if additional_info:
docstring += set_min_indent(
f"{param_name} ({param_type}{additional_info}):{description}",
indent_level + 8,
)
else:
docstring += set_min_indent(
f"{param_name} ({param_type}{shape_string}{optional_string}{param_default}):{description}",
indent_level + 8,
)
# TODO (Yoni): Add support for Attributes section in docs
else:
print(
f"You used `@auto_class_docstring` decorator on `{cls.__name__}` but this class is not part of the AutoMappings. Remove the decorator"
)
# Assign the dynamically generated docstring to the wrapper class
cls.__doc__ = docstring
return cls