pared down model

model.py

@@ -1,16 +1,16 @@
 import torch
 import lightning
-from torch.utils.data import Dataset
-from typing import Any, Dict
-import argparse
+# from torch.utils.data import Dataset
+# from typing import Any, Dict
+# import argparse
 from pydantic import BaseModel
-from get_dataset_dictionaries import get_dict_pair
-import os
-import shutil
+# from get_dataset_dictionaries import get_dict_pair
+# import os
+# import shutil
 
-import optuna
-from optuna.integration import PyTorchLightningPruningCallback
-from functools import partial
+# import optuna
+# from optuna.integration import PyTorchLightningPruningCallback
+# from functools import partial
 
 class FFNModule(torch.nn.Module):
     """
@@ -104,295 +104,295 @@ class FeatureNormPredictor(lightning.LightningModule):
         self.model.load_state_dict(torch.load(path))
 
     
-class HiddenStateFeatureNormDataset(Dataset):
-    def __init__(
-        self, 
-        input_embeddings: Dict[str, torch.Tensor],
-        feature_norms: Dict[str, torch.Tensor],
-    ):
+# class HiddenStateFeatureNormDataset(Dataset):
+#     def __init__(
+#         self, 
+#         input_embeddings: Dict[str, torch.Tensor],
+#         feature_norms: Dict[str, torch.Tensor],
+#     ):
         
-        # Invariant: input_embeddings and target_feature_norms have exactly the same keys
-        # this should be done by the train/test split and upstream data processing
-        assert(input_embeddings.keys() == feature_norms.keys())
-
-        self.words = list(input_embeddings.keys())
-        self.input_embeddings = torch.stack([
-            input_embeddings[word] for word in self.words
-        ])
-        self.feature_norms = torch.stack([
-            feature_norms[word] for word in self.words
-        ])
+#         # Invariant: input_embeddings and target_feature_norms have exactly the same keys
+#         # this should be done by the train/test split and upstream data processing
+#         assert(input_embeddings.keys() == feature_norms.keys())
+
+#         self.words = list(input_embeddings.keys())
+#         self.input_embeddings = torch.stack([
+#             input_embeddings[word] for word in self.words
+#         ])
+#         self.feature_norms = torch.stack([
+#             feature_norms[word] for word in self.words
+#         ])
         
-    def __len__(self):
-        return len(self.words)
+#     def __len__(self):
+#         return len(self.words)
     
-    def __getitem__(self, idx):
-        return self.input_embeddings[idx], self.feature_norms[idx]
-
-# this is used when not optimizing
-def train(args : Dict[str, Any]):
-
-    # input_embeddings = torch.load(args.input_embeddings)
-    # feature_norms = torch.load(args.feature_norms)
-    # words = list(input_embeddings.keys())
-
-    input_embeddings, feature_norms, norm_list = get_dict_pair(
-        args.norm,
-        args.embedding_dir,
-        args.lm_layer,
-        translated= False if args.raw_buchanan else True,
-        normalized= True if args.normal_buchanan else False
-    ) 
-    norms_file = open(args.save_dir+"/"+args.save_model_name+'.txt','w')
-    norms_file.write("\n".join(norm_list))
-    norms_file.close()
-
-    words = list(input_embeddings.keys())
+#     def __getitem__(self, idx):
+#         return self.input_embeddings[idx], self.feature_norms[idx]
+
+# # this is used when not optimizing
+# def train(args : Dict[str, Any]):
+
+#     # input_embeddings = torch.load(args.input_embeddings)
+#     # feature_norms = torch.load(args.feature_norms)
+#     # words = list(input_embeddings.keys())
+
+#     input_embeddings, feature_norms, norm_list = get_dict_pair(
+#         args.norm,
+#         args.embedding_dir,
+#         args.lm_layer,
+#         translated= False if args.raw_buchanan else True,
+#         normalized= True if args.normal_buchanan else False
+#     ) 
+#     norms_file = open(args.save_dir+"/"+args.save_model_name+'.txt','w')
+#     norms_file.write("\n".join(norm_list))
+#     norms_file.close()
+
+#     words = list(input_embeddings.keys())
     
-    model = FeatureNormPredictor(
-        FFNParams(
-            input_size=input_embeddings[words[0]].shape[0],
-            output_size=feature_norms[words[0]].shape[0],
-            hidden_size=args.hidden_size,
-            num_layers=args.num_layers,
-            dropout=args.dropout,
-        ),
-        TrainingParams(
-            num_epochs=args.num_epochs,
-            batch_size=args.batch_size,
-            learning_rate=args.learning_rate,
-            weight_decay=args.weight_decay,
-        ),
-    )
-
-    # train/val split
-    train_size = int(len(words) * 0.8)
-    valid_size = len(words) - train_size
-    train_words, validation_words = torch.utils.data.random_split(words, [train_size, valid_size])
-
-    # TODO: Methodology Decision: should we be normalizing the hidden states/feature norms?
-    train_embeddings = {word: input_embeddings[word] for word in train_words}
-    train_feature_norms = {word: feature_norms[word] for word in train_words}
-    validation_embeddings = {word: input_embeddings[word] for word in validation_words}
-    validation_feature_norms = {word: feature_norms[word] for word in validation_words}
-
-    train_dataset = HiddenStateFeatureNormDataset(train_embeddings, train_feature_norms)
-    train_dataloader = torch.utils.data.DataLoader(
-        train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-    )
-    validation_dataset = HiddenStateFeatureNormDataset(validation_embeddings, validation_feature_norms)
-    validation_dataloader = torch.utils.data.DataLoader(
-        validation_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-    )
-
-    callbacks = [
-        lightning.pytorch.callbacks.ModelCheckpoint(
-            save_last=True,
-            dirpath=args.save_dir,
-            filename=args.save_model_name,
-        ),
-    ]
-    if args.early_stopping is not None:
-        callbacks.append(lightning.pytorch.callbacks.EarlyStopping(
-            monitor="val_loss",
-            patience=args.early_stopping,
-            mode='min',
-            min_delta=0.0
-        ))
-
-    #TODO Design Decision - other trainer args? Is device necessary?
-    # cpu is fine for the scale of this model - only a few layers and a few hundred words
-    trainer = lightning.Trainer(
-        max_epochs=args.num_epochs,
-        callbacks=callbacks,
-        accelerator="cpu",
-        log_every_n_steps=7
-    )
-
-    trainer.fit(model, train_dataloader, validation_dataloader)
-
-    trainer.validate(model, validation_dataloader)
-
-    return model
-
-# this is used when optimizing
-def objective(trial: optuna.trial.Trial, args: Dict[str, Any]) -> float:
-    # optimizing hidden size, batch size, and learning rate
-    input_embeddings, feature_norms, norm_list = get_dict_pair(
-        args.norm,
-        args.embedding_dir,
-        args.lm_layer,
-        translated= False if args.raw_buchanan else True,
-        normalized= True if args.normal_buchanan else False
-    )
-    norms_file = open(args.save_dir+"/"+args.save_model_name+'.txt','w')
-    norms_file.write("\n".join(norm_list))
-    norms_file.close()
-
-    words = list(input_embeddings.keys())
-    input_size=input_embeddings[words[0]].shape[0]
-    output_size=feature_norms[words[0]].shape[0]
-    min_size = min(output_size, input_size)
-    max_size = min(output_size, 2*input_size)if min_size == input_size else min(2*output_size, input_size)
-    hidden_size = trial.suggest_int("hidden_size", min_size, max_size, log=True)
-    batch_size = trial.suggest_int("batch_size", 16, 128, log=True)
-    learning_rate = trial.suggest_float("learning_rate", 1e-6, 1, log=True)
-
-    model = FeatureNormPredictor(
-        FFNParams(
-            input_size=input_size,
-            output_size=output_size,
-            hidden_size=hidden_size,
-            num_layers=args.num_layers,
-            dropout=args.dropout,
-        ),
-        TrainingParams(
-            num_epochs=args.num_epochs,
-            batch_size=batch_size,
-            learning_rate=learning_rate,
-            weight_decay=args.weight_decay,
-        ),
-    )
-
-    # train/val split
-    train_size = int(len(words) * 0.8)
-    valid_size = len(words) - train_size
-    train_words, validation_words = torch.utils.data.random_split(words, [train_size, valid_size])
-
-    train_embeddings = {word: input_embeddings[word] for word in train_words}
-    train_feature_norms = {word: feature_norms[word] for word in train_words}
-    validation_embeddings = {word: input_embeddings[word] for word in validation_words}
-    validation_feature_norms = {word: feature_norms[word] for word in validation_words}
-
-    train_dataset = HiddenStateFeatureNormDataset(train_embeddings, train_feature_norms)
-    train_dataloader = torch.utils.data.DataLoader(
-        train_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-    )
-    validation_dataset = HiddenStateFeatureNormDataset(validation_embeddings, validation_feature_norms)
-    validation_dataloader = torch.utils.data.DataLoader(
-        validation_dataset,
-        batch_size=args.batch_size,
-        shuffle=True,
-    )
-
-    callbacks = [
-        # all trial models will be saved in temporary directory
-        lightning.pytorch.callbacks.ModelCheckpoint(
-            save_last=True,
-            dirpath=os.path.join(args.save_dir,'optuna_trials'),
-            filename="{}".format(trial.number)
-        ),
-    ]
-
-    if args.prune is not None:
-        callbacks.append(PyTorchLightningPruningCallback(
-            trial,
-            monitor='val_loss'
-        ))
+#     model = FeatureNormPredictor(
+#         FFNParams(
+#             input_size=input_embeddings[words[0]].shape[0],
+#             output_size=feature_norms[words[0]].shape[0],
+#             hidden_size=args.hidden_size,
+#             num_layers=args.num_layers,
+#             dropout=args.dropout,
+#         ),
+#         TrainingParams(
+#             num_epochs=args.num_epochs,
+#             batch_size=args.batch_size,
+#             learning_rate=args.learning_rate,
+#             weight_decay=args.weight_decay,
+#         ),
+#     )
+
+#     # train/val split
+#     train_size = int(len(words) * 0.8)
+#     valid_size = len(words) - train_size
+#     train_words, validation_words = torch.utils.data.random_split(words, [train_size, valid_size])
+
+#     # TODO: Methodology Decision: should we be normalizing the hidden states/feature norms?
+#     train_embeddings = {word: input_embeddings[word] for word in train_words}
+#     train_feature_norms = {word: feature_norms[word] for word in train_words}
+#     validation_embeddings = {word: input_embeddings[word] for word in validation_words}
+#     validation_feature_norms = {word: feature_norms[word] for word in validation_words}
+
+#     train_dataset = HiddenStateFeatureNormDataset(train_embeddings, train_feature_norms)
+#     train_dataloader = torch.utils.data.DataLoader(
+#         train_dataset,
+#         batch_size=args.batch_size,
+#         shuffle=True,
+#     )
+#     validation_dataset = HiddenStateFeatureNormDataset(validation_embeddings, validation_feature_norms)
+#     validation_dataloader = torch.utils.data.DataLoader(
+#         validation_dataset,
+#         batch_size=args.batch_size,
+#         shuffle=True,
+#     )
+
+#     callbacks = [
+#         lightning.pytorch.callbacks.ModelCheckpoint(
+#             save_last=True,
+#             dirpath=args.save_dir,
+#             filename=args.save_model_name,
+#         ),
+#     ]
+#     if args.early_stopping is not None:
+#         callbacks.append(lightning.pytorch.callbacks.EarlyStopping(
+#             monitor="val_loss",
+#             patience=args.early_stopping,
+#             mode='min',
+#             min_delta=0.0
+#         ))
+
+#     #TODO Design Decision - other trainer args? Is device necessary?
+#     # cpu is fine for the scale of this model - only a few layers and a few hundred words
+#     trainer = lightning.Trainer(
+#         max_epochs=args.num_epochs,
+#         callbacks=callbacks,
+#         accelerator="cpu",
+#         log_every_n_steps=7
+#     )
+
+#     trainer.fit(model, train_dataloader, validation_dataloader)
+
+#     trainer.validate(model, validation_dataloader)
+
+#     return model
+
+# # this is used when optimizing
+# def objective(trial: optuna.trial.Trial, args: Dict[str, Any]) -> float:
+#     # optimizing hidden size, batch size, and learning rate
+#     input_embeddings, feature_norms, norm_list = get_dict_pair(
+#         args.norm,
+#         args.embedding_dir,
+#         args.lm_layer,
+#         translated= False if args.raw_buchanan else True,
+#         normalized= True if args.normal_buchanan else False
+#     )
+#     norms_file = open(args.save_dir+"/"+args.save_model_name+'.txt','w')
+#     norms_file.write("\n".join(norm_list))
+#     norms_file.close()
+
+#     words = list(input_embeddings.keys())
+#     input_size=input_embeddings[words[0]].shape[0]
+#     output_size=feature_norms[words[0]].shape[0]
+#     min_size = min(output_size, input_size)
+#     max_size = min(output_size, 2*input_size)if min_size == input_size else min(2*output_size, input_size)
+#     hidden_size = trial.suggest_int("hidden_size", min_size, max_size, log=True)
+#     batch_size = trial.suggest_int("batch_size", 16, 128, log=True)
+#     learning_rate = trial.suggest_float("learning_rate", 1e-6, 1, log=True)
+
+#     model = FeatureNormPredictor(
+#         FFNParams(
+#             input_size=input_size,
+#             output_size=output_size,
+#             hidden_size=hidden_size,
+#             num_layers=args.num_layers,
+#             dropout=args.dropout,
+#         ),
+#         TrainingParams(
+#             num_epochs=args.num_epochs,
+#             batch_size=batch_size,
+#             learning_rate=learning_rate,
+#             weight_decay=args.weight_decay,
+#         ),
+#     )
+
+#     # train/val split
+#     train_size = int(len(words) * 0.8)
+#     valid_size = len(words) - train_size
+#     train_words, validation_words = torch.utils.data.random_split(words, [train_size, valid_size])
+
+#     train_embeddings = {word: input_embeddings[word] for word in train_words}
+#     train_feature_norms = {word: feature_norms[word] for word in train_words}
+#     validation_embeddings = {word: input_embeddings[word] for word in validation_words}
+#     validation_feature_norms = {word: feature_norms[word] for word in validation_words}
+
+#     train_dataset = HiddenStateFeatureNormDataset(train_embeddings, train_feature_norms)
+#     train_dataloader = torch.utils.data.DataLoader(
+#         train_dataset,
+#         batch_size=args.batch_size,
+#         shuffle=True,
+#     )
+#     validation_dataset = HiddenStateFeatureNormDataset(validation_embeddings, validation_feature_norms)
+#     validation_dataloader = torch.utils.data.DataLoader(
+#         validation_dataset,
+#         batch_size=args.batch_size,
+#         shuffle=True,
+#     )
+
+#     callbacks = [
+#         # all trial models will be saved in temporary directory
+#         lightning.pytorch.callbacks.ModelCheckpoint(
+#             save_last=True,
+#             dirpath=os.path.join(args.save_dir,'optuna_trials'),
+#             filename="{}".format(trial.number)
+#         ),
+#     ]
+
+#     if args.prune is not None:
+#         callbacks.append(PyTorchLightningPruningCallback(
+#             trial,
+#             monitor='val_loss'
+#         ))
         
-    if args.early_stopping is not None:
-        callbacks.append(lightning.pytorch.callbacks.EarlyStopping(
-            monitor="val_loss",
-            patience=args.early_stopping,
-            mode='min',
-            min_delta=0.0
-        ))
-    # note that if optimizing is chosen, will automatically not implement vanilla early stopping 
-    #TODO Design Decision - other trainer args? Is device necessary?
-    # cpu is fine for the scale of this model - only a few layers and a few hundred words
-    trainer = lightning.Trainer(
-        max_epochs=args.num_epochs,
-        callbacks=callbacks,
-        accelerator="cpu",
-        log_every_n_steps=7,
-        # enable_checkpointing=False
-    )
-
-    trainer.fit(model, train_dataloader, validation_dataloader)
-
-    trainer.validate(model, validation_dataloader)
+#     if args.early_stopping is not None:
+#         callbacks.append(lightning.pytorch.callbacks.EarlyStopping(
+#             monitor="val_loss",
+#             patience=args.early_stopping,
+#             mode='min',
+#             min_delta=0.0
+#         ))
+#     # note that if optimizing is chosen, will automatically not implement vanilla early stopping 
+#     #TODO Design Decision - other trainer args? Is device necessary?
+#     # cpu is fine for the scale of this model - only a few layers and a few hundred words
+#     trainer = lightning.Trainer(
+#         max_epochs=args.num_epochs,
+#         callbacks=callbacks,
+#         accelerator="cpu",
+#         log_every_n_steps=7,
+#         # enable_checkpointing=False
+#     )
+
+#     trainer.fit(model, train_dataloader, validation_dataloader)
+
+#     trainer.validate(model, validation_dataloader)
     
-    return trainer.callback_metrics['val_loss'].item()
-
-if __name__ == "__main__":
-    # parse args
-    parser = argparse.ArgumentParser()
-    #TODO: Design Decision: Should we input paths, to the pre-extracted layers, or the model/layer we want to generate them from
-    # required inputs
-    parser.add_argument("--norm", type=str, required=True, help="feature norm set to use")
-    parser.add_argument("--embedding_dir", type=str, required=True, help=" directory containing embeddings")
-    parser.add_argument("--lm_layer", type=int, required=True, help="layer of embeddings to use")
-    # if user selects optimize, hidden_size, batch_size and learning_rate will be optimized. 
-    parser.add_argument("--optimize", action="store_true", help="optimize hyperparameters for training")
-    parser.add_argument("--prune", action="store_true", help="prune unpromising trials when optimizing")
-    # optional hyperparameter specs
-    parser.add_argument("--num_layers", type=int, default=2, help="number of layers in FFN")
-    parser.add_argument("--hidden_size", type=int, default=100, help="hidden size of FFN")
-    parser.add_argument("--dropout", type=float, default=0.1, help="dropout rate of FFN")
-    # set this to at least 100 if doing early stopping
-    parser.add_argument("--num_epochs", type=int, default=10, help="number of epochs to train for")
-    parser.add_argument("--batch_size", type=int, default=32, help="batch size for training")
-    parser.add_argument("--learning_rate", type=float, default=0.001, help="learning rate for training")
-    parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for training")
-    parser.add_argument("--early_stopping", type=int, default=None, help="number of epochs to wait for early stopping")
-    # optional dataset specs, for buchanan really
-    parser.add_argument('--raw_buchanan', action="store_true", help="do not use translated values for buchanan")
-    parser.add_argument('--normal_buchanan', action="store_true", help="use normalized features for buchanan")
-    # required for output
-    parser.add_argument("--save_dir", type=str, required=True, help="directory to save model to")
-    parser.add_argument("--save_model_name", type=str, required=True, help="name of model to save")
-
-    args = parser.parse_args()
-
-    if args.early_stopping is not None:
-        args.num_epochs = max(50, args.num_epochs)
-
-    torch.manual_seed(10)
-
-    if args.optimize:
-        # call optimizer code here
-        print("optimizing for learning rate, batch size, and hidden size")
-        pruner = optuna.pruners.MedianPruner() if args.prune else optuna.pruners.NopPruner()
-        sampler = optuna.samplers.TPESampler(seed=10)
-
-        study = optuna.create_study(direction='minimize', pruner=pruner, sampler=sampler)
-        study.optimize(partial(objective, args=args), n_trials = 100, timeout=600)
-
-        other_params = {
-            "num_layers": args.num_layers,
-            "num_epochs": args.num_epochs,
-            "dropout": args.dropout,
-            "weight_decay": args.weight_decay,
-        }
-
-        print("Number of finished trials: {}".format(len(study.trials)))
-
-        trial = study.best_trial
-        print("Best trial: "+str(trial.number))
+#     return trainer.callback_metrics['val_loss'].item()
+
+# if __name__ == "__main__":
+#     # parse args
+#     parser = argparse.ArgumentParser()
+#     #TODO: Design Decision: Should we input paths, to the pre-extracted layers, or the model/layer we want to generate them from
+#     # required inputs
+#     parser.add_argument("--norm", type=str, required=True, help="feature norm set to use")
+#     parser.add_argument("--embedding_dir", type=str, required=True, help=" directory containing embeddings")
+#     parser.add_argument("--lm_layer", type=int, required=True, help="layer of embeddings to use")
+#     # if user selects optimize, hidden_size, batch_size and learning_rate will be optimized. 
+#     parser.add_argument("--optimize", action="store_true", help="optimize hyperparameters for training")
+#     parser.add_argument("--prune", action="store_true", help="prune unpromising trials when optimizing")
+#     # optional hyperparameter specs
+#     parser.add_argument("--num_layers", type=int, default=2, help="number of layers in FFN")
+#     parser.add_argument("--hidden_size", type=int, default=100, help="hidden size of FFN")
+#     parser.add_argument("--dropout", type=float, default=0.1, help="dropout rate of FFN")
+#     # set this to at least 100 if doing early stopping
+#     parser.add_argument("--num_epochs", type=int, default=10, help="number of epochs to train for")
+#     parser.add_argument("--batch_size", type=int, default=32, help="batch size for training")
+#     parser.add_argument("--learning_rate", type=float, default=0.001, help="learning rate for training")
+#     parser.add_argument("--weight_decay", type=float, default=0.0, help="weight decay for training")
+#     parser.add_argument("--early_stopping", type=int, default=None, help="number of epochs to wait for early stopping")
+#     # optional dataset specs, for buchanan really
+#     parser.add_argument('--raw_buchanan', action="store_true", help="do not use translated values for buchanan")
+#     parser.add_argument('--normal_buchanan', action="store_true", help="use normalized features for buchanan")
+#     # required for output
+#     parser.add_argument("--save_dir", type=str, required=True, help="directory to save model to")
+#     parser.add_argument("--save_model_name", type=str, required=True, help="name of model to save")
+
+#     args = parser.parse_args()
+
+#     if args.early_stopping is not None:
+#         args.num_epochs = max(50, args.num_epochs)
+
+#     torch.manual_seed(10)
+
+#     if args.optimize:
+#         # call optimizer code here
+#         print("optimizing for learning rate, batch size, and hidden size")
+#         pruner = optuna.pruners.MedianPruner() if args.prune else optuna.pruners.NopPruner()
+#         sampler = optuna.samplers.TPESampler(seed=10)
+
+#         study = optuna.create_study(direction='minimize', pruner=pruner, sampler=sampler)
+#         study.optimize(partial(objective, args=args), n_trials = 100, timeout=600)
+
+#         other_params = {
+#             "num_layers": args.num_layers,
+#             "num_epochs": args.num_epochs,
+#             "dropout": args.dropout,
+#             "weight_decay": args.weight_decay,
+#         }
+
+#         print("Number of finished trials: {}".format(len(study.trials)))
+
+#         trial = study.best_trial
+#         print("Best trial: "+str(trial.number))
         
 
-        print("  Validation Loss: {}".format(trial.value))
+#         print("  Validation Loss: {}".format(trial.value))
 
-        print("  Optimized Params: ")
-        for key, value in trial.params.items():
-            print("    {}: {}".format(key, value))
+#         print("  Optimized Params: ")
+#         for key, value in trial.params.items():
+#             print("    {}: {}".format(key, value))
 
-        print("  User Defined Params: ")
-        for key, value in other_params.items():
-            print("    {}: {}".format(key, value))
+#         print("  User Defined Params: ")
+#         for key, value in other_params.items():
+#             print("    {}: {}".format(key, value))
         
-        print('saving best trial')
-        for filename in os.listdir(os.path.join(args.save_dir,'optuna_trials')):
-            if filename == "{}.ckpt".format(trial.number):
-                shutil.move(os.path.join(args.save_dir,'optuna_trials',filename), os.path.join(args.save_dir, "{}.ckpt".format(args.save_model_name)))
-        shutil.rmtree(os.path.join(args.save_dir,'optuna_trials'))
-
-    else:   
-        model = train(args)
+#         print('saving best trial')
+#         for filename in os.listdir(os.path.join(args.save_dir,'optuna_trials')):
+#             if filename == "{}.ckpt".format(trial.number):
+#                 shutil.move(os.path.join(args.save_dir,'optuna_trials',filename), os.path.join(args.save_dir, "{}.ckpt".format(args.save_model_name)))
+#         shutil.rmtree(os.path.join(args.save_dir,'optuna_trials'))
+
+#     else:   
+#         model = train(args)
         

requirements.txt

@@ -1,4 +1,6 @@
 torch
 minicons
 pandas
-numpy
+numpy
+pydantic
+lightning