Upload weakly_supervised_parser/utils/populate_chart.py

weakly_supervised_parser/utils/populate_chart.py

@@ -26,9 +26,9 @@ 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
+from pytorch_lightning import Trainer
 
-# trainer = Trainer(accelerator="auto", enable_progress_bar=False, max_epochs=-1)
+trainer = Trainer(accelerator="auto", enable_progress_bar=False, max_epochs=-1)
 
 
 class PopulateCKYChart:
@@ -54,20 +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])
+            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[