SimpleSentimentModel should inherit from BatchedModel.

This class implements the abstract method `predict_minibatch` in `BatchedModel`.

Addressing #1361.

PiperOrigin-RevId: 586041904

lit_nlp/examples/simple_pytorch_demo.py

@@ -79,7 +79,7 @@ def _from_pretrained(cls, *args, **kw):
     return cls.from_pretrained(*args, from_tf=True, **kw)
 
 
-class SimpleSentimentModel(lit_model.Model):
+class SimpleSentimentModel(lit_model.BatchedModel):
   """Simple sentiment analysis model."""
 
   LABELS = ["0", "1"]  # negative, positive
@@ -103,7 +103,7 @@ def __init__(self, model_name_or_path):
   ##
   # LIT API implementation
   def max_minibatch_size(self):
-    # This tells lit_model.Model.predict() how to batch inputs to
+    # This tells lit_model.BatchedModel.predict() how to batch inputs to
     # predict_minibatch().
     # Alternately, you can just override predict() and handle batching yourself.
     return 32

lit_nlp/examples/sst_pytorch_demo.py

@@ -70,7 +70,7 @@ def _from_pretrained(cls, *args, **kw):
     return cls.from_pretrained(*args, from_tf=True, **kw)
 
 
-class SimpleSentimentModel(lit_model.Model):
+class SimpleSentimentModel(lit_model.BatchedModel):
   """Simple sentiment analysis model."""
 
   LABELS = ["0", "1"]  # negative, positive
@@ -95,7 +95,7 @@ def __init__(self, model_name_or_path):
   ##
   # LIT API implementation
   def max_minibatch_size(self):
-    # This tells lit_model.Model.predict() how to batch inputs to
+    # This tells lit_model.BatchedModel.predict() how to batch inputs to
     # predict_minibatch().
     # Alternately, you can just override predict() and handle batching yourself.
     return 32