Files
2026-07-13 12:49:22 +08:00

136 lines
4.7 KiB
Python

"""
Module that handles the string representation of Datalab objects.
"""
from abc import ABC, abstractmethod
from typing import TYPE_CHECKING, Dict, List, Optional, Type
from cleanlab.datalab.internal.task import Task
if TYPE_CHECKING: # pragma: no cover
from cleanlab.datalab.internal.data_issues import DataIssues
class RepresentationStrategy(ABC):
def __init__(self, data_issues: "DataIssues"):
self.data_issues = data_issues
@property
def checks_run(self) -> bool:
return not self.data_issues.issues.empty
@property
def num_examples(self) -> Optional[int]:
return self.data_issues.get_info("statistics").get("num_examples")
@property
def num_classes(self) -> Optional[int]:
return self.data_issues.get_info("statistics").get("num_classes")
@property
def issues_identified(self) -> str:
return (
self.data_issues.issue_summary["num_issues"].sum() if self.checks_run else "Not checked"
)
def show_task(self, task: "Task") -> str:
return f"task={str(task).capitalize()}"
def show_checks_run(self) -> str:
return f"checks_run={self.checks_run}"
def show_num_examples(self) -> str:
return f"num_examples={self.num_examples}" if self.num_examples is not None else ""
def show_num_classes(self) -> str:
return f"num_classes={self.num_classes}" if self.num_classes is not None else ""
def show_issues_identified(self) -> str:
return f"issues_identified={self.issues_identified}"
@abstractmethod
def represent(self) -> str:
pass
def to_string(self, task: "Task") -> str:
"""What is displayed if user executes: print(datalab)"""
info_list = [
f"Task: {str(task).capitalize()}",
f"Checks run: {'Yes' if self.checks_run else 'No'}",
f"Number of examples: {self.num_examples if self.num_examples is not None else 'Unknown'}",
f"Number of classes: {self.num_classes if self.num_classes is not None else 'Unknown'}",
f"Issues identified: {self.issues_identified}",
]
return "Datalab:\n" + "\n".join(info_list)
class ClassificationRepresentation(RepresentationStrategy):
def represent(self) -> str:
display_strings: List[str] = [
self.show_task(Task.CLASSIFICATION),
self.show_checks_run(),
self.show_num_examples(),
self.show_num_classes(),
self.show_issues_identified(),
]
# Drop empty strings
display_strings = [s for s in display_strings if bool(s)]
display_str = ", ".join(display_strings)
return f"Datalab({display_str})"
class RegressionRepresentation(RepresentationStrategy):
def represent(self) -> str:
display_strings: List[str] = [
self.show_task(Task.REGRESSION),
self.show_checks_run(),
self.show_num_examples(),
self.show_issues_identified(),
]
# Drop empty strings
display_strings = [s for s in display_strings if bool(s)]
display_str = ", ".join(display_strings)
return f"Datalab({display_str})"
class MultilabelRepresentation(RepresentationStrategy):
def represent(self) -> str:
display_strings: List[str] = [
self.show_task(Task.MULTILABEL),
self.show_checks_run(),
self.show_num_examples(),
self.show_num_classes(),
self.show_issues_identified(),
]
# Drop empty strings
display_strings = [s for s in display_strings if bool(s)]
display_str = ", ".join(display_strings)
return f"Datalab({display_str})"
class _Displayer:
def __init__(self, data_issues: "DataIssues", task: "Task") -> None:
self.data_issues = data_issues
self.task = task
self.representation_strategy = self._get_representation_strategy()
def _get_representation_strategy(self) -> RepresentationStrategy:
strategies: Dict[str, Type[RepresentationStrategy]] = {
"classification": ClassificationRepresentation,
"regression": RegressionRepresentation,
"multilabel": MultilabelRepresentation,
}
strategy_class = strategies.get(self.task.value)
if strategy_class is None:
raise ValueError(f"Unsupported task type: {self.task}")
return strategy_class(self.data_issues)
def __repr__(self) -> str:
"""What is displayed in console if user executes: >>> datalab"""
return self.representation_strategy.represent()
def __str__(self) -> str:
"""What is displayed if user executes: print(datalab)"""
return self.representation_strategy.to_string(self.task)