update model ckpt

app.py

@@ -10,14 +10,16 @@ from weakly_supervised_parser.tree.evaluate import calculate_F1_for_spans, tree_
 from weakly_supervised_parser.inference import Predictor
 from weakly_supervised_parser.model.trainer import InsideOutsideStringClassifier
 
+from weakly_supervised_parser.model.span_classifier import LightningModel
+
 benepar.download('benepar_en3')
 
 nlp = spacy.load("en_core_web_md")
 nlp.add_pipe("benepar", config={"model": "benepar_en3"})
 
-inside_model = InsideOutsideStringClassifier(model_name_or_path="roberta-base", max_seq_length=256)
-fetch_url_inside_model = hf_hub_url(repo_id="nickil/weakly-supervised-parsing", filename="inside_model.onnx", revision="main")
-inside_model.load_model(pre_trained_model_path=cached_download(fetch_url_inside_model))
+# inside_model = InsideOutsideStringClassifier(model_name_or_path="roberta-base", max_seq_length=256)
+fetch_url_inside_model = hf_hub_url(repo_id="nickil/weakly-supervised-parsing", filename="inside_model.ckpt", revision="main")
+inside_model = LightningModel.load_from_checkpoint(checkpoint_path=cached_download(fetch_url_inside_model))
 
 # outside_model = InsideOutsideStringClassifier(model_name_or_path="roberta-base", max_seq_length=64)
 # outside_model.load_model(pre_trained_model_path=TRAINED_MODEL_PATH + "outside_model.onnx")
@@ -35,7 +37,7 @@ def predict(sentence, model):
     elif model == "inside-outside":
         best_parse = Predictor(sentence=sentence).obtain_best_parse(predict_type="inside_outside", model=inside_outside_model, scale_axis=1, predict_batch_size=128)
     sentence_f1 = calculate_F1_for_spans(tree_to_spans(gold_standard), tree_to_spans(best_parse))
-    return gold_standard, best_parse, sentence_f1
+    return gold_standard, best_parse, f"{sentence_f1:.2f}"
 
 
 iface = gradio.Interface(

weakly_supervised_parser/inference.py

@@ -8,6 +8,8 @@ from weakly_supervised_parser.tree.evaluate import calculate_F1_for_spans, tree_
 from weakly_supervised_parser.model.trainer import InsideOutsideStringClassifier
 from weakly_supervised_parser.settings import PTB_TEST_SENTENCES_WITHOUT_PUNCTUATION_PATH, PTB_TEST_GOLD_WITHOUT_PUNCTUATION_ALIGNED_PATH
 
+from weakly_supervised_parser.model.span_classifier import LightningModel
+
 
 class Predictor:
     def __init__(self, sentence):
@@ -96,7 +98,7 @@ def main():
     args = parser.parse_args()
 
     if args.use_inside:
-        pre_trained_model_path = TRAINED_MODEL_PATH + "inside_model.onnx"
+        pre_trained_model_path = TRAINED_MODEL_PATH + "inside_model.ckpt"
         max_seq_length = args.inside_max_seq_length
 
     if args.use_inside_self_train:
@@ -116,8 +118,10 @@ def main():
         outside_model = InsideOutsideStringClassifier(model_name_or_path=args.model_name_or_path, max_seq_length=args.outside_max_seq_length)
         outside_model.load_model(pre_trained_model_path=outside_pre_trained_model_path)
     else:
-        model = InsideOutsideStringClassifier(model_name_or_path=args.model_name_or_path, max_seq_length=max_seq_length)
-        model.load_model(pre_trained_model_path=pre_trained_model_path)
+        # model = InsideOutsideStringClassifier(model_name_or_path=args.model_name_or_path, max_seq_length=max_seq_length)
+        # model.load_model(pre_trained_model_path=pre_trained_model_path)
+        
+        model = LightningModel.load_from_checkpoint(checkpoint_path=pre_trained_model_path)
 
     if args.use_inside or args.use_inside_self_train:
         predict_type = "inside"
@@ -126,7 +130,6 @@ def main():
         predict_type = "outside"
 
     with open(args.save_path, "w") as out_file:
-        print(type(args.scale_axis))
         test_sentences = DataLoaderHelper(input_file_object=PTB_TEST_SENTENCES_WITHOUT_PUNCTUATION_PATH).read_lines()
         test_gold_file_path = PTB_TEST_GOLD_WITHOUT_PUNCTUATION_ALIGNED_PATH
         for test_index, test_sentence in enumerate(test_sentences):

weakly_supervised_parser/model/data_module_loader.py

@@ -35,13 +35,15 @@ class PyTorchDataModule(Dataset):
             add_special_tokens=True,
             return_tensors="pt",
         )
-
+        
         out = dict(
             sentence=sentence,
             input_ids=sentence_encoding["input_ids"].flatten(),
             attention_mask=sentence_encoding["attention_mask"].flatten(),
-            labels=data_row["label"].flatten(),
         )
+        
+        if "label" in self.data.columns:
+            out.update(dict(labels=data_row["label"].flatten()))
 
         return out
 
@@ -52,6 +54,7 @@ class DataModule(LightningDataModule):
         model_name_or_path: str,
         train_df: pd.DataFrame,
         eval_df: pd.DataFrame,
+        test_df: pd.DataFrame,
         max_seq_length: int = 256,
         train_batch_size: int = 32,
         eval_batch_size: int = 32,
@@ -62,6 +65,7 @@ class DataModule(LightningDataModule):
         self.model_name_or_path = model_name_or_path
         self.train_df = train_df
         self.eval_df = eval_df
+        self.test_df = test_df
         self.max_seq_length = max_seq_length
         self.train_batch_size = train_batch_size
         self.eval_batch_size = eval_batch_size
@@ -71,9 +75,15 @@ class DataModule(LightningDataModule):
 
         self.train_dataset = PyTorchDataModule(self.model_name_or_path, self.train_df, self.max_seq_length)
         self.eval_dataset = PyTorchDataModule(self.model_name_or_path, self.eval_df, self.max_seq_length)
+        
+        if isinstance(self.test_df, pd.DataFrame):
+            self.test_dataset = PyTorchDataModule(self.model_name_or_path, self.test_df, self.max_seq_length)
 
     def train_dataloader(self) -> DataLoader:
         return DataLoader(self.train_dataset, batch_size=self.train_batch_size, shuffle=False, num_workers=self.num_workers, pin_memory=True)
 
     def val_dataloader(self) -> DataLoader:
         return DataLoader(self.eval_dataset, batch_size=self.eval_batch_size, shuffle=False, num_workers=self.num_workers, pin_memory=True)
+
+    def predict_dataloader(self) -> DataLoader:
+        return DataLoader(self.test_dataset, batch_size=len(self.test_dataset), shuffle=False, num_workers=self.num_workers, pin_memory=True)

weakly_supervised_parser/model/span_classifier.py

@@ -45,6 +45,11 @@ class LightningModel(LightningModule):
         preds = torch.argmax(logits, axis=1)
         labels = batch["labels"]
         return {"loss": val_loss, "preds": preds, "labels": labels}
+    
+    def predict_step(self, batch, batch_idx, dataloader_idx=0):
+        batch = {"input_ids": batch["input_ids"], "attention_mask": batch["attention_mask"]}
+        outputs = self(**batch)
+        return torch.nn.functional.softmax(outputs.logits, dim=1)[:, 1]
 
     def validation_epoch_end(self, outputs):
         preds = torch.cat([x["preds"] for x in outputs])
@@ -58,6 +63,8 @@ class LightningModel(LightningModule):
         return loss
 
     def setup(self, stage=None):
+        if stage != "fit":
+            return None
         # Get dataloader by calling it - train_dataloader() is called after setup() by default
         train_loader = self.trainer.datamodule.train_dataloader()
 

weakly_supervised_parser/model/trainer.py

@@ -10,7 +10,7 @@ from pytorch_lightning import Trainer, seed_everything
 from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
 from transformers import AutoTokenizer, logging
 
-from onnxruntime import InferenceSession, SessionOptions
+from onnxruntime import InferenceSession
 from scipy.special import softmax
 
 from weakly_supervised_parser.model.data_module_loader import DataModule
@@ -37,7 +37,7 @@ class InsideOutsideStringClassifier:
         devices: int = 1,
         enable_progress_bar: bool = True,
         enable_model_summary: bool = False,
-        enable_checkpointing: bool = False,
+        enable_checkpointing: bool = True,
         logger: bool = False,
         accelerator: str = "auto",
         train_batch_size: int = 32,
@@ -52,6 +52,7 @@ class InsideOutsideStringClassifier:
             model_name_or_path=self.model_name_or_path,
             train_df=train_df,
             eval_df=eval_df,
+            test_df=None,
             max_seq_length=self.max_seq_length,
             train_batch_size=train_batch_size,
             eval_batch_size=eval_batch_size,
@@ -70,7 +71,7 @@ class InsideOutsideStringClassifier:
 
         callbacks = []
         callbacks.append(EarlyStopping(monitor="val_loss", patience=2, mode="min", check_finite=True))
-        # callbacks.append(ModelCheckpoint(monitor="val_loss", dirpath=outputdir, filename=filename, save_top_k=1, save_weights_only=True, mode="min"))
+        callbacks.append(ModelCheckpoint(monitor="val_loss", dirpath=outputdir, filename=filename, save_top_k=1, save_weights_only=True, mode="min"))
 
         trainer = Trainer(
             accelerator=accelerator,
@@ -98,10 +99,7 @@ class InsideOutsideStringClassifier:
         )
 
     def load_model(self, pre_trained_model_path):
-        options = SessionOptions()
-        options.intra_op_num_threads = 32
-        options.inter_op_num_threads = 32
-        self.model = InferenceSession(pre_trained_model_path, options, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
+        self.model = InferenceSession(pre_trained_model_path, providers=["CUDAExecutionProvider", "CPUExecutionProvider"])
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_name_or_path, use_fast=True)
 
     def preprocess_function(self, data):
@@ -129,3 +127,24 @@ class InsideOutsideStringClassifier:
 
     def predict(self, spans):
         return self.predict_proba(spans).argmax(axis=1)
+
+    
+class InsideOutsideStringPredictor:
+
+    def __init__(self, model_name_or_path, max_seq_length, pre_trained_model_path, num_workers=32):
+        self.model_name_or_path = model_name_or_path
+        self.pre_trained_model_path = pre_trained_model_path
+        self.max_seq_length = max_seq_length
+        self.num_workers = num_workers 
+
+    def predict_proba(self, test_df):
+        test_dataloader = data_module = DataModule(
+                                        model_name_or_path=self.model_name_or_path,
+                                        train_df=None,
+                                        eval_df=None,
+                                        test_df=test_df,
+                                        max_seq_length=self.max_seq_length,
+                                        num_workers=self.num_workers,
+                                    )
+
+        return trainer.predict(model, dataloaders=test_dataloader)

weakly_supervised_parser/utils/populate_chart.py

@@ -1,17 +1,24 @@
 import pandas as pd
 import numpy as np
+import logging
 
 from datasets.utils import set_progress_bar_enabled
 
 from weakly_supervised_parser.utils.prepare_dataset import NGramify
 from weakly_supervised_parser.utils.create_inside_outside_strings import InsideOutside
+from weakly_supervised_parser.model.trainer import InsideOutsideStringPredictor
 from weakly_supervised_parser.utils.cky_algorithm import get_best_parse
 from weakly_supervised_parser.utils.distant_supervision import RuleBasedHeuristic
 from weakly_supervised_parser.utils.prepare_dataset import PTBDataset
 from weakly_supervised_parser.settings import PTB_TRAIN_SENTENCES_WITHOUT_PUNCTUATION_PATH
 
+from weakly_supervised_parser.model.data_module_loader import DataModule
+from weakly_supervised_parser.model.span_classifier import LightningModel
+
 # Disable Dataset.map progress bar
 set_progress_bar_enabled(False)
+logging.getLogger("pytorch_lightning").setLevel(logging.WARNING)
+
 
 # ptb = PTBDataset(data_path=PTB_TRAIN_SENTENCES_WITHOUT_PUNCTUATION_PATH)
 # ptb_top_100_common = [item.lower() for item in RuleBasedHeuristic(corpus=ptb.retrieve_all_sentences()).get_top_tokens(top_most_common_ptb=100)]
@@ -19,6 +26,10 @@ ptb_top_100_common = ['this', 'myself', 'shouldn', 'not', 'analysts', 'same', 'm
 # ptb_most_common_first_token = RuleBasedHeuristic(corpus=ptb.retrieve_all_sentences()).augment_using_most_frequent_starting_token(N=1)[0][0].lower()
 ptb_most_common_first_token = "the"
 
+from pytorch_lightning import Trainer
+
+trainer = Trainer(accelerator="auto", enable_progress_bar=False, gpus=-1)
+
 
 class PopulateCKYChart:
     def __init__(self, sentence):
@@ -43,17 +54,20 @@ class PopulateCKYChart:
 
         if predict_type == "inside":
             
-            if data.shape[0] > chunks:
-                data_chunks = np.array_split(data, data.shape[0] // chunks)
-                for data_chunk in data_chunks:
-                    inside_scores.extend(model.predict_proba(spans=data_chunk.rename(columns={"inside_sentence": "sentence"})[["sentence"]],
-                                                             scale_axis=scale_axis,
-                                                             predict_batch_size=predict_batch_size)[:, 1])
-            else:
-                inside_scores.extend(model.predict_proba(spans=data.rename(columns={"inside_sentence": "sentence"})[["sentence"]],
-                                                         scale_axis=scale_axis,
-                                                         predict_batch_size=predict_batch_size)[:, 1])
-
+#             if data.shape[0] > chunks:
+#                 data_chunks = np.array_split(data, data.shape[0] // chunks)
+#                 for data_chunk in data_chunks:
+#                     inside_scores.extend(model.predict_proba(spans=data_chunk.rename(columns={"inside_sentence": "sentence"})[["sentence"]],
+#                                                              scale_axis=scale_axis,
+#                                                              predict_batch_size=predict_batch_size)[:, 1])
+#             else:
+#                 inside_scores.extend(model.predict_proba(spans=data.rename(columns={"inside_sentence": "sentence"})[["sentence"]],
+#                                                          scale_axis=scale_axis,
+#                                                          predict_batch_size=predict_batch_size)[:, 1])
+            test_dataloader = DataModule(model_name_or_path="roberta-base", train_df=None, eval_df=None, 
+                                         test_df=data.rename(columns={"inside_sentence": "sentence"})[["sentence"]])
+            inside_scores.extend(trainer.predict(model, dataloaders=test_dataloader)[0])
+                                 
             data["inside_scores"] = inside_scores
             data.loc[
                 (data["inside_sentence"].str.lower().str.startswith(ptb_most_common_first_token))