Pretty-printing of Model objects

PiperOrigin-RevId: 629571604
PAIR-code · bdu91 · Jun 26, 2024 · Apr 19, 2024 · Apr 26, 2024 · Apr 26, 2024
commit 4fb3bde897c68fdeb3bd829f6e5a88223bc131a4
diff --git a/lit_nlp/api/model.py b/lit_nlp/api/model.py
@@ -89,6 +89,18 @@ def description(self) -> str:
     """
     return inspect.getdoc(self) or ''
 
+  def __str__(self) -> str:
+    classname = self.__class__.__module__ + '.' + self.__class__.__qualname__
+    indented_description = '  ' + self.description().replace('\n', '\n  ')
+    return f'{classname}(...):\n{indented_description}'
+
+  def _repr_pretty_(self, p, cycle):
+    """Pretty-printing for IPython environments, both notebooks and repl."""
+    if not cycle:
+      p.text(str(self))
+    else:
+      p.text('...')
+
   @classmethod
   def init_spec(cls) -> Optional[Spec]:
     """Attempts to infer a Spec describing a Model's constructor parameters.