From 4fb3bde897c68fdeb3bd829f6e5a88223bc131a4 Mon Sep 17 00:00:00 2001
From: Ian Tenney <iftenney@google.com>
Date: Tue, 30 Apr 2024 17:09:01 -0700
Subject: [PATCH] Pretty-printing of Model objects

PiperOrigin-RevId: 629571604
---
 lit_nlp/api/model.py | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/lit_nlp/api/model.py b/lit_nlp/api/model.py
index ce034a20..49a68649 100644
--- 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.