enums.py 3.6 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153
  1. # Copyright The Lightning team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. from lightning_utilities.core.enums import StrEnum
  15. from typing_extensions import Literal
  16. class EnumStr(StrEnum):
  17. """Base Enum."""
  18. @staticmethod
  19. def _name() -> str:
  20. return "Task"
  21. @classmethod
  22. def from_str(cls: type["EnumStr"], value: str, source: Literal["key", "value", "any"] = "key") -> "EnumStr":
  23. """Load from string.
  24. Raises:
  25. ValueError:
  26. If required value is not among the supported options.
  27. >>> class MyEnum(EnumStr):
  28. ... a = "aaa"
  29. ... b = "bbb"
  30. >>> MyEnum.from_str("a")
  31. <MyEnum.a: 'aaa'>
  32. >>> MyEnum.from_str("c")
  33. Traceback (most recent call last):
  34. ...
  35. ValueError: Invalid Task: expected one of ['a', 'b'], but got c.
  36. """
  37. try:
  38. me = super().from_str(value.replace("-", "_"), source=source)
  39. except ValueError as err:
  40. _allowed_im = [m.lower() for m in cls._member_names_]
  41. raise ValueError(
  42. f"Invalid {cls._name()}: expected one of {cls._allowed_matches(source)}, but got {value}."
  43. ) from err
  44. return cls(me)
  45. class DataType(EnumStr):
  46. """Enum to represent data type.
  47. >>> "Binary" in list(DataType)
  48. True
  49. """
  50. @staticmethod
  51. def _name() -> str:
  52. return "Data type"
  53. BINARY = "binary"
  54. MULTILABEL = "multi-label"
  55. MULTICLASS = "multi-class"
  56. MULTIDIM_MULTICLASS = "multi-dim multi-class"
  57. class AverageMethod(EnumStr):
  58. """Enum to represent average method.
  59. >>> None in list(AverageMethod)
  60. True
  61. >>> AverageMethod.NONE == None
  62. True
  63. >>> AverageMethod.NONE == 'none'
  64. True
  65. """
  66. @staticmethod
  67. def _name() -> str:
  68. return "Average method"
  69. MICRO = "micro"
  70. MACRO = "macro"
  71. WEIGHTED = "weighted"
  72. NONE = None
  73. SAMPLES = "samples"
  74. class MDMCAverageMethod(EnumStr):
  75. """Enum to represent multi-dim multi-class average method."""
  76. @staticmethod
  77. def _name() -> str:
  78. return "MDMC Average method"
  79. GLOBAL = "global"
  80. SAMPLEWISE = "samplewise"
  81. class ClassificationTask(EnumStr):
  82. """Enum to represent the different tasks in classification metrics.
  83. >>> "binary" in list(ClassificationTask)
  84. True
  85. """
  86. @staticmethod
  87. def _name() -> str:
  88. return "Classification"
  89. BINARY = "binary"
  90. MULTICLASS = "multiclass"
  91. MULTILABEL = "multilabel"
  92. class ClassificationTaskNoBinary(EnumStr):
  93. """Enum to represent the different tasks in classification metrics.
  94. >>> "binary" in list(ClassificationTaskNoBinary)
  95. False
  96. """
  97. @staticmethod
  98. def _name() -> str:
  99. return "Classification"
  100. MULTILABEL = "multilabel"
  101. MULTICLASS = "multiclass"
  102. class ClassificationTaskNoMultilabel(EnumStr):
  103. """Enum to represent the different tasks in classification metrics.
  104. >>> "multilabel" in list(ClassificationTaskNoMultilabel)
  105. False
  106. """
  107. @staticmethod
  108. def _name() -> str:
  109. return "Classification"
  110. BINARY = "binary"
  111. MULTICLASS = "multiclass"