Enable batch processing in scriptable tokenizer example

examples/text_classification_with_scriptable_tokenizer/handler.py

@@ -1,6 +1,5 @@
 """
 Module for text classification with scriptable tokenizer
-DOES NOT SUPPORT BATCH!
 """
 import logging
 from abc import ABC
@@ -51,18 +50,19 @@ def preprocess(self, data):
 
         # Compat layer: normally the envelope should just return the data
         # directly, but older versions of Torchserve didn't have envelope.
-        # Processing only the first input, not handling batch inference
 
-        line = data[0]
-        text = line.get("data") or line.get("body")
-        # Decode text if not a str but bytes or bytearray
-        if isinstance(text, (bytes, bytearray)):
-            text = text.decode("utf-8")
+        text_batch = []
+        for line in data:
+            text = line.get("data") or line.get("body")
+            # Decode text if not a str but bytes or bytearray
+            if isinstance(text, (bytes, bytearray)):
+                text = text.decode("utf-8")
 
-        text = remove_html_tags(text)
-        text = text.lower()
+            text = remove_html_tags(text)
+            text = text.lower()
+            text_batch.append(text)
 
-        return text
+        return text_batch
 
     def inference(self, data, *args, **kwargs):
         """The Inference Request is made through this function and the user

examples/text_classification_with_scriptable_tokenizer/script_tokenizer_and_model.py

@@ -76,7 +76,7 @@ def main(args):
     model = XLMR_BASE_ENCODER.get_model(head=classifier_head)
 
     # Load trained parameters and load them into the model
-    model.load_state_dict(torch.load(args.input_file))
+    model.load_state_dict(torch.load(args.input_file, map_location=torch.device("cpu")))
 
     # Chain the tokenizer, the adapter and the model
     combi_model = T.Sequential(
@@ -88,7 +88,7 @@ def main(args):
     combi_model.eval()
 
     # Make sure to move the model to CPU to avoid placement error during loading
-    combi_model.to("cpu")
+    combi_model.to(torch.device("cpu"))
 
     combi_model_jit = torch.jit.script(combi_model)