Pretty-printing of Model objects

iftenney · LIT team · commit 4fb3bde897c6 · 2024-04-30T17:10:21.000-07:00
PiperOrigin-RevId: 629571604
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.
