[#1] main_infer.py implemented

idiomify/fetchers.py

@@ -95,8 +95,7 @@ def fetch_alpha(ver: str, run: Run = None) -> Alpha:
     artifact_dir = artifact.download(root=alpha_dir(ver))
     ckpt_path = path.join(artifact_dir, "model.ckpt")
     bart = AutoModelForSeq2SeqLM.from_config(AutoConfig.from_pretrained(config['bart']))
-    with open(ckpt_path, 'r') as fh:
-        alpha = Alpha.load_from_checkpoint(ckpt_path, bart=bart)
+    alpha = Alpha.load_from_checkpoint(ckpt_path, bart=bart)
     return alpha
 
 

idiomify/idiomifier.py

@@ -0,0 +1,22 @@
+from transformers import BartTokenizer
+from builders import SourcesBuilder
+from models import Alpha
+
+
+class Idiomifier:
+
+    def __init__(self, model: Alpha, tokenizer: BartTokenizer):
+        self.model = model
+        self.builder = SourcesBuilder(tokenizer)
+        self.model.eval()
+
+    def __call__(self, src: str, max_length=100) -> str:
+        srcs = self.builder(literal2idiomatic=[(src, "")])
+        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

idiomify/models.py

@@ -47,14 +47,6 @@ class Alpha(pl.LightningModule):  # noqa
     def on_train_batch_end(self, outputs: dict, *args, **kwargs):
         self.log("Train/Loss", outputs['loss'])
 
-    def predict(self, srcs: torch.Tensor) -> torch.Tensor:
-        pred_ids = self.bart.generate(
-            inputs=srcs[:, 0],  # (N, 2, L) -> (N, L)
-            attention_mask=srcs[:, 1],  # (N, 2, L) -> (N, L)
-            decoder_start_token_id=self.hparams['bos_token_id'],
-        )
-        return pred_ids  # (N, L)
-
     def configure_optimizers(self) -> torch.optim.Optimizer:
         """
         Instantiates and returns the optimizer to be used for this model

main_infer.py

@@ -1,37 +1,28 @@
-# we disable them for now.
-# import argparse
-# from idiomify.fetchers import fetch_config, fetch_rd, fetch_idioms
-# from transformers import BertTokenizer
-# from termcolor import colored
-#
-#
-# def main():
-#         parser = argparse.ArgumentParser()
-#         parser.add_argument("--model", type=str,
-#                             default="alpha")
-#         parser.add_argument("--ver", type=str,
-#                             default="eng2eng")
-#         parser.add_argument("--sent", type=str,
-#                             default="to avoid getting to the point")
-#         args = parser.parse_args()
-#         config = fetch_config()[args.model][args.ver]
-#         config.update(vars(args))
-#         idioms = fetch_idioms(config['idioms_ver'])
-#         rd = fetch_rd(config['model'], config['ver'])
-#         rd.eval()
-#         tokenizer = BertTokenizer.from_pretrained(config['bert'])
-#         X = T.inputs([("", config['sent'])], tokenizer, config['k'])
-#         probs = rd.P_wisdom(X).squeeze().tolist()
-#         wisdom2prob = [
-#                 (wisdom, prob)
-#                 for wisdom, prob in zip(idioms, probs)
-#         ]
-#         # sort and append
-#         res = list(sorted(wisdom2prob, key=lambda x: x[1], reverse=True))
-#         print(f"query: {colored(text=config['sent'], color='blue')}")
-#         for idx, (idiom, prob) in enumerate(res):
-#             print(idx, idiom, prob)
-#
-#
-# if __name__ == '__main__':
-#     main()
+import argparse
+from termcolor import colored
+from idiomifier import Idiomifier
+from idiomify.fetchers import fetch_config, fetch_alpha
+from transformers import BartTokenizer
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str,
+                        default="alpha")
+    parser.add_argument("--ver", type=str,
+                        default="overfit")
+    parser.add_argument("--src", type=str,
+                        default="If there's any benefits to losing 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()
+    config = fetch_config()[args.model][args.ver]
+    config.update(vars(args))
+    model = fetch_alpha(config['ver'])
+    tokenizer = BartTokenizer.from_pretrained(config['bart'])
+    idiomifier = Idiomifier(model, tokenizer)
+    src = config['src']
+    tgt = idiomifier(src=config['src'])
+    print(src, "\n->", colored(tgt, "blue"))
+
+
+if __name__ == '__main__':
+    main()