[#2] Pipeline now supports multiple sentences with batch_decode

config.yaml

@@ -7,6 +7,7 @@ idiomifier:
   max_epochs: 2
   batch_size: 40
   shuffle: true
+  seed: 104
 
 # for building & uploading datasets or tokenizer
 idioms:

idiomify/pipeline.py

@@ -1,5 +1,7 @@
 
 # for inference
+from typing import List
+
 from transformers import BartTokenizer
 
 from builders import SourcesBuilder
@@ -12,13 +14,13 @@ class Pipeline:
         self.model = model
         self.builder = SourcesBuilder(tokenizer)
 
-    def __call__(self, src: str, max_length=100) -> str:
-        srcs = self.builder(literal2idiomatic=[(src, "")])
+    def __call__(self, sents: List[str], max_length=100) -> List[str]:
+        srcs = self.builder(literal2idiomatic=[(sent, "") for sent in sents])
         pred_ids = self.model.bart.generate(
             inputs=srcs[:, 0],  # (N, 2, L) -> (N, L)
             attention_mask=srcs[:, 1],  # (N, 2, L) -> (N, L)
             decoder_start_token_id=self.model.hparams['bos_token_id'],
             max_length=max_length,
-        ).squeeze()  # -> (N, L_t) -> (L_t)
-        tgt = self.builder.tokenizer.decode(pred_ids, skip_special_tokens=True)
-        return tgt
+        )  # -> (N, L_t)
+        tgts = self.builder.tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
+        return tgts

main_infer.py

@@ -1,12 +1,15 @@
+"""
+This is for just a simple sanity check on the inference.
+"""
 import argparse
-from idiomify.models import Pipeline
+from idiomify.pipeline import Pipeline
 from idiomify.fetchers import fetch_config, fetch_idiomifier
 from transformers import BartTokenizer
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--src", type=str,
+    parser.add_argument("--sent", type=str,
                         default="If there's any good to loosing my job,"
                                 " it's that I'll now be able to go to school full-time and finish my degree earlier.")
     args = parser.parse_args()
@@ -16,8 +19,8 @@ def main():
     model.eval()  # this is crucial
     tokenizer = BartTokenizer.from_pretrained(config['bart'])
     pipeline = Pipeline(model, tokenizer)
-    src = config['src']
-    tgt = pipeline(src=config['src'])
+    src = config['sent']
+    tgt = pipeline(sents=[config['sent']])
     print(src, "\n->", tgt)
 
 

main_train.py

@@ -23,7 +23,6 @@ def main():
     config.update(vars(args))
     if not config['upload']:
         print(colored("WARNING: YOU CHOSE NOT TO UPLOAD. NOTHING BUT LOGS WILL BE SAVED TO WANDB", color="red"))
-
     # prepare the model
     bart = BartForConditionalGeneration.from_pretrained(config['bart'])
     tokenizer = BartTokenizer.from_pretrained(config['bart'])