Repository: asuprem/EdnaML (language: python)

Issue:
minor changes for dependabot



Discussion / hints:

Fix patch:
diff --git a/src/ednaml/crawlers/CUB200_2011DataCrawler.py b/archives/old crawlers/CUB200_2011DataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/CUB200_2011DataCrawler.py
rename to archives/old crawlers/CUB200_2011DataCrawler.py
diff --git a/src/ednaml/crawlers/Cars196DataCrawler.py b/archives/old crawlers/Cars196DataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/Cars196DataCrawler.py
rename to archives/old crawlers/Cars196DataCrawler.py
diff --git a/src/ednaml/crawlers/CoLabelCompCarsCrawler.py b/archives/old crawlers/CoLabelCompCarsCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/CoLabelCompCarsCrawler.py
rename to archives/old crawlers/CoLabelCompCarsCrawler.py
diff --git a/src/ednaml/crawlers/CoLabelIntegratedDatasetCrawler.py b/archives/old crawlers/CoLabelIntegratedDatasetCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/CoLabelIntegratedDatasetCrawler.py
rename to archives/old crawlers/CoLabelIntegratedDatasetCrawler.py
diff --git a/src/ednaml/crawlers/MTMCDataCrawler.py b/archives/old crawlers/MTMCDataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/MTMCDataCrawler.py
rename to archives/old crawlers/MTMCDataCrawler.py
diff --git a/src/ednaml/crawlers/Market1501DataCrawler.py b/archives/old crawlers/Market1501DataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/Market1501DataCrawler.py
rename to archives/old crawlers/Market1501DataCrawler.py
diff --git a/src/ednaml/crawlers/SUNDataCrawler.py b/archives/old crawlers/SUNDataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/SUNDataCrawler.py
rename to archives/old crawlers/SUNDataCrawler.py
diff --git a/src/ednaml/crawlers/VRICDataCrawler.py b/archives/old crawlers/VRICDataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/VRICDataCrawler.py
rename to archives/old crawlers/VRICDataCrawler.py
diff --git a/src/ednaml/crawlers/VeRiDataCrawler.py b/archives/old crawlers/VeRiDataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/VeRiDataCrawler.py
rename to archives/old crawlers/VeRiDataCrawler.py
diff --git a/src/ednaml/crawlers/VehicleColorCrawler.py b/archives/old crawlers/VehicleColorCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/VehicleColorCrawler.py
rename to archives/old crawlers/VehicleColorCrawler.py
diff --git a/src/ednaml/crawlers/VehicleIDDataCrawler.py b/archives/old crawlers/VehicleIDDataCrawler.py
similarity index 100%
rename from src/ednaml/crawlers/VehicleIDDataCrawler.py
rename to archives/old crawlers/VehicleIDDataCrawler.py
diff --git a/src/ednaml/datareaders/CompCars.py b/archives/old crawlers/old datareaders/CompCars.py
similarity index 100%
rename from src/ednaml/datareaders/CompCars.py
rename to archives/old crawlers/old datareaders/CompCars.py
diff --git a/src/ednaml/datareaders/VeRi.py b/archives/old crawlers/old datareaders/VeRi.py
similarity index 100%
rename from src/ednaml/datareaders/VeRi.py
rename to archives/old crawlers/old datareaders/VeRi.py
diff --git a/src/ednaml/datareaders/VehicleColor.py b/archives/old crawlers/old datareaders/VehicleColor.py
similarity index 100%
rename from src/ednaml/datareaders/VehicleColor.py
rename to archives/old crawlers/old datareaders/VehicleColor.py
diff --git a/src/ednaml/datareaders/VehicleID.py b/archives/old crawlers/old datareaders/VehicleID.py
similarity index 100%
rename from src/ednaml/datareaders/VehicleID.py
rename to archives/old crawlers/old datareaders/VehicleID.py
diff --git a/archives/requirements.txt b/archives/old_requirements.txt
similarity index 100%
rename from archives/requirements.txt
rename to archives/old_requirements.txt
diff --git a/profiles/FNC/experts/configs/config_base.yml b/profiles/FNC/experts/configs/config_base.yml
new file mode 100644
index 00000000..4f619751
--- /dev/null
+++ b/profiles/FNC/experts/configs/config_base.yml
@@ -0,0 +1,38 @@
+EXECUTION:
+  EPOCHS: 10
+  FP16: False
+  MODEL_SERVING: False
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  SKIPEVAL: False
+  TEST_FREQUENCY: 1
+  TRAINER: HFTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 8
+    
+LOSS: 
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    LABEL: fnews
+    NAME: classification
+
+OPTIMIZER:
+  - OPTIMIZER: AdamW
+    BASE_LR: 1.0e-5
+    LR_BIAS_FACTOR: 1.0
+    OPTIMIZER_KWARGS: 
+      eps: 1.0e-6
+    OPTIMIZER_NAME: opt-1
+    WEIGHT_BIAS_FACTOR: 0.0005
+    WEIGHT_DECAY: 0.0005
+
+SCHEDULER:
+  - LR_KWARGS: 
+      step_size: 5
+      gamma: 0.5
+    LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
diff --git a/profiles/FNC/experts/configs/config_base_plugin_deploy.yml b/profiles/FNC/experts/configs/config_base_plugin_deploy.yml
new file mode 100644
index 00000000..66f123a5
--- /dev/null
+++ b/profiles/FNC/experts/configs/config_base_plugin_deploy.yml
@@ -0,0 +1,12 @@
+DEPLOYMENT:
+  DEPLOY: HFDeploy
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: warmup
+    RESET: False
+
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
diff --git a/profiles/FNC/experts/configs/crawler-efnd-base.yml b/profiles/FNC/experts/configs/crawler-efnd-base.yml
new file mode 100644
index 00000000..eff6ebf3
--- /dev/null
+++ b/profiles/FNC/experts/configs/crawler-efnd-base.yml
@@ -0,0 +1,64 @@
+EXECUTION:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data
+      include: ["cmu_miscov19"] # Adjust this to specific EFND dataset
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-efnd-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+DEPLOYMENT:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data
+      include: ["cmu_miscov19"] # Adjust this to specific EFND dataset
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-efnd-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-expert-efnd
+  MODEL_BACKBONE: unknown
+  MODEL_QUALIFIER: cmu_miscov19
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 3 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
\ No newline at end of file
diff --git a/profiles/FNC/experts/configs/crawler-fakeddit-base.yml b/profiles/FNC/experts/configs/crawler-fakeddit-base.yml
new file mode 100644
index 00000000..95cbf412
--- /dev/null
+++ b/profiles/FNC/experts/configs/crawler-fakeddit-base.yml
@@ -0,0 +1,62 @@
+EXECUTION:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data/Fakeddit
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-fakeddit-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+DEPLOYMENT:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data/Fakeddit
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-fakeddit-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+      
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-expert-fakeddit
+  MODEL_BACKBONE: unknown
+  MODEL_QUALIFIER: fakeddit
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 3 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
\ No newline at end of file
diff --git a/profiles/FNC/experts/configs/crawler-nela-base.yml b/profiles/FNC/experts/configs/crawler-nela-base.yml
new file mode 100644
index 00000000..2048b819
--- /dev/null
+++ b/profiles/FNC/experts/configs/crawler-nela-base.yml
@@ -0,0 +1,66 @@
+EXECUTION:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data
+      sub_folder: nela-covid-2020
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-nela-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: ["covid", "corona", "mask", "wuhan", "n95", "sars", "monkey", "pandemic", "social", "quarantin", "virus", "infect", "lock", "ppe", "variant", "vaccine", "travel", "omicron", "ivermectin", 
+        "plandemic", "5g", "gates", "hoax", "bioweapon", "bat", "fauci"]
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+DEPLOYMENT:
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      data_folder: Data
+      sub_folder: nela-covid-2020
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-nela-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-expert-nela
+  MODEL_BACKBONE: unknown
+  MODEL_QUALIFIER: nela_covid_2020
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 3 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
\ No newline at end of file
diff --git a/profiles/FNC/experts/configs/model-bertvariant.yml b/profiles/FNC/experts/configs/model-bertvariant.yml
new file mode 100644
index 00000000..741d1564
--- /dev/null
+++ b/profiles/FNC/experts/configs/model-bertvariant.yml
@@ -0,0 +1,26 @@
+EXECUTION:
+  DATAREADER:
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+
+DEPLOYMENT:
+  DATAREADER:
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: HFAutoModel
+  MODEL_BASE: albert-base-v2
+  MODEL_KWARGS: 
+    auto_class: AutoModelForSequenceClassification
+    hidden_act: gelu
+    pooling: pooled
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
+
+SAVE:
+  MODEL_BACKBONE: albert
+  MODEL_VERSION: 2  # to deal with other bugs...
\ No newline at end of file
diff --git a/profiles/FNC/experts/configs/plugin_base.yml b/profiles/FNC/experts/configs/plugin_base.yml
new file mode 100644
index 00000000..90a9e1b9
--- /dev/null
+++ b/profiles/FNC/experts/configs/plugin_base.yml
@@ -0,0 +1,39 @@
+DEPLOYMENT:
+  EPOCHS: 2
+  
+MODEL_PLUGIN:
+  - PLUGIN: FastKMeansProxy
+    PLUGIN_NAME: FastKMP-l2
+    PLUGIN_KWARGS:
+      proxies: 20
+      dimensions: 768
+      dist: euclidean
+      iterations: 30
+      batch_size: 256
+      alpha: 0.6
+      classifier_access: encoder.classifier
+  - PLUGIN: FastKMeansProxy
+    PLUGIN_NAME: FastKMP-cos
+    PLUGIN_KWARGS:
+      proxies: 20
+      dimensions: 768
+      dist: cosine
+      iterations: 30
+      batch_size: 256
+      alpha: 0.6
+      classifier_access: encoder.classifier
+  - PLUGIN: FastRandomizedLipschitz
+    PLUGIN_NAME: FRL-midas
+    PLUGIN_KWARGS:
+      proxies: 15
+      dimensions: 768
+      dist: euclidean
+      proxy_epochs: 1
+      iterations: 30
+      batch_size: 256
+      alpha: 0.6
+      classifier_access: encoder.classifier
+  - PLUGIN: LogitConfidence
+    PLUGIN_NAME: logit-confidence
+    PLUGIN_KWARGS:
+      num_classes: 2
\ No newline at end of file
diff --git a/profiles/FNC/experts/crawlers/efnd-crawler.py b/profiles/FNC/experts/crawlers/efnd-crawler.py
new file mode 100644
index 00000000..f9183a70
--- /dev/null
+++ b/profiles/FNC/experts/crawlers/efnd-crawler.py
@@ -0,0 +1,63 @@
+import ednaml
+import csv, glob, os
+from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
+
+@edna.register_crawler
+class MiDASCrawler(Crawler):
+  def __init__(self, logger = None, data_folder="Data", include=["cmu_miscov19", "covid19_fn_title", "kagglefn_short"]):
+    """Crawls the Data folder with all datasets already extracted to their individual folders.
+    Assumes specific file construction: files are separated into splits 
+    (train, test, and val), with label subsets fake and true, with naming convention:
+
+
+    Crawl yields lists with tuples of:
+    <text>, <shares>, <likes>, <label>, <label...s>
+
+    """
+    logger.info("Crawling %s for %s"%(data_folder, str(include)))
+    
+    self.classes = {}
+    self.classes["fnews"] = 2
+    self.metadata = {}
+
+    self.data_folder = data_folder
+    datasets_folders = glob.glob(os.path.join(data_folder, "*"))
+    datasets_folders = [item for item in datasets_folders if os.path.basename(item) in include]
+    self.classes["dataset"] = len(datasets_folders)
+
+
+    self.metadata["train"] = {}
+    self.metadata["test"] = {}
+    self.metadata["val"] = {}
+    self.metadata["train"]["crawl"] = []
+    self.metadata["test"]["crawl"] = []
+    self.metadata["val"]["crawl"] = []
+    for idx, folder in enumerate(datasets_folders):
+      ftrain = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "train.csv"])) for subset in ["fake", "true"]]
+      ftest = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "test.csv"])) for subset in ["fake", "true"]]
+      fval = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "val.csv"])) for subset in ["fake", "true"]]
+
+      # train:
+      self.metadata["train"]["crawl"] += self.getTextAndLabels(ftrain, idx)
+      self.metadata["test"]["crawl"] += self.getTextAndLabels(ftest, idx)
+      self.metadata["val"]["crawl"] += self.getTextAndLabels(fval, idx)
+
+    self.metadata["train"]["classes"] = self.classes
+    self.metadata["test"]["classes"] = self.classes
+    self.metadata["val"]["classes"] = self.classes
+
+  def getTextAndLabels(self, listOfFiles, datasetidx):
+    crawl = []
+    for file in listOfFiles:
+      # For each file, open as csv. Extract the relevant columns ("text", "label")
+      with open(file, "r") as ofile:
+        csvread = csv.reader(ofile)
+
+        header = next(csvread)
+        textidx = [idx for idx,item in enumerate(header) if item == "text"][0]
+        labelidx = [idx for idx,item in enumerate(header) if item == "label"][0]
+
+        for row in csvread:
+          crawl.append([row[textidx], 0,0,int(row[labelidx]), datasetidx])
+    return crawl
diff --git a/profiles/FNC/experts/crawlers/fakeddit-crawler.py b/profiles/FNC/experts/crawlers/fakeddit-crawler.py
new file mode 100644
index 00000000..5d6a365e
--- /dev/null
+++ b/profiles/FNC/experts/crawlers/fakeddit-crawler.py
@@ -0,0 +1,63 @@
+import ednaml, torch, os, csv
+from ednaml.crawlers import Crawler
+
+import ednaml.core.decorators as edna
+
+@edna.register_crawler
+class FakedditCrawler(Crawler):
+    def __init__(self, logger = None, data_folder="Data/Fakeddit"):
+        """Crawls the Data folder with all datasets already extracted to their individual folders.
+        Assumes specific file construction: files are separated into splits 
+        (train, test, and val), with label subsets fake and true, with naming convention:
+
+        Crawl yields lists with tuples of:
+        <text>, <shares>, <likes>, <label>, <label...s>
+
+        """
+        logger.info("Crawling %s"%(data_folder))
+        
+        # set up class metadata
+        self.classes = {}
+        self.classes["fnews"] = 2
+
+        # set up paths
+        self.data_folder = data_folder
+        self.trainfile = os.path.join(data_folder, "all_train.tsv")
+        self.valfile = os.path.join(data_folder, "all_validate.tsv")
+        self.testfile = os.path.join(data_folder, "all_test_public.tsv")
+
+        # set up content metadata
+        self.metadata = {}
+        self.metadata["train"] = {}
+        self.metadata["test"] = {}
+        self.metadata["val"] = {}
+        self.metadata["train"]["crawl"] = []
+        self.metadata["test"]["crawl"] = []
+        self.metadata["val"]["crawl"] = []
+        
+        # Row information for Fakeddit:
+        #   -- 5:  clean text
+        #   -- 13:  score (likes)
+        #   -- 17:  2 way label
+        #   -- -1:  6-way label
+        #   -- 12: subreddit name
+        with open(self.trainfile, "r") as fl:
+            fobj = csv.reader(fl, delimiter="\t")
+            header = next(fobj)
+            for row in fobj:
+                self.metadata["train"]["crawl"] += [(row[5], 0, int(row[13]), int(row[17]), int(row[-1]))]
+        with open(self.valfile, "r") as fl:
+            fobj = csv.reader(fl, delimiter="\t")
+            header = next(fobj)
+            for row in fobj:
+                self.metadata["val"]["crawl"] += [(row[5], 0, int(row[13]), int(row[17]), int(row[-1]))]
+        with open(self.testfile, "r") as fl:
+            fobj = csv.reader(fl, delimiter="\t")
+            header = next(fobj)
+            for row in fobj:
+                self.metadata["test"]["crawl"] += [(row[5], 0, int(row[13]), int(row[17]), int(row[-1]))]
+        
+
+        self.metadata["train"]["classes"] = self.classes
+        self.metadata["test"]["classes"] = self.classes
+        self.metadata["val"]["classes"] = self.classes
\ No newline at end of file
diff --git a/profiles/FNC/experts/crawlers/nela-crawler.py b/profiles/FNC/experts/crawlers/nela-crawler.py
new file mode 100644
index 00000000..d92c03b7
--- /dev/null
+++ b/profiles/FNC/experts/crawlers/nela-crawler.py
@@ -0,0 +1,101 @@
+import ednaml, torch, os, csv, os, json, glob, click
+from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
+
+
+@edna.register_crawler
+class NELACrawler(Crawler):
+  def __init__(self, logger = None, data_folder="Data", sub_folder="nela-covid-2020"):
+    """Crawls the Data folder with all datasets already extracted to their individual folders.
+    Assumes specific file construction: files are separated into splits 
+    (train, test, and val), with label subsets fake and true, with naming convention:
+
+    Crawl yields lists with tuples of:
+    <text>, <shares>, <likes>, <label>, <label...s>
+
+    """
+    logger.info("Crawling %s"%(data_folder))
+    
+    # set up class metadata
+    self.classes = {}
+    self.classes["fnews"] = 2
+    
+    # set up paths
+    self.data_folder = os.path.join(data_folder , sub_folder)
+    labelfile = os.path.join(self.data_folder, "labels.csv")
+    self.data_folder = os.path.join(self.data_folder, sub_folder) # because there is nested
+    ndata = "newsdata"
+    tweet = "tweet"
+    
+    self.newsdata = os.path.join(self.data_folder, ndata)
+    self.tweetdata = os.path.join(self.data_folder, tweet)
+
+    # obtain source labels
+    sourcelabels = {}
+    with open(labelfile, "r") as lfile:
+      labelobj = csv.reader(lfile)
+      header = next(labelobj)
+      for row in labelobj:
+        sourcelabels[row[0]] = int(row[1])
+
+    # set up content metadata
+    self.metadata = {}
+    self.metadata["train"] = {}
+    self.metadata["test"] = {}
+    self.metadata["val"] = {}
+    self.metadata["train"]["crawl"] = []
+    self.metadata["test"]["crawl"] = []
+    self.metadata["val"]["crawl"] = []
+    
+    
+    # We don't need the tweet-data because it is the same as the newsdata
+    #with open(os.path.join(self.tweetdata, "tweet.json")) as obj:
+    #  self.tweetdata = json.load(obj)
+
+
+    # for each file in newsdata: obtain the titles, and save in label propagated. We willa dd to metadata later...
+    newsdataitems = glob.glob(os.path.join(self.newsdata, "*.json"))
+    labelsets = {0:[],1:[],2:[]}
+    for newsdatafile in newsdataitems:
+      newssource = os.path.splitext(os.path.basename(newsdatafile))[0]
+      label = sourcelabels.get(newssource,1)
+      njson = None
+      with open(newsdatafile) as jsonobj:
+        njson = json.load(jsonobj)
+      labelsets[label] += [item["title"] for item in njson]
+
+    # shuffle
+    import random
+    random.seed(75837)
+    random.shuffle(labelsets[0])
+    random.seed(75837)
+    random.shuffle(labelsets[1])
+    random.seed(75837)
+    random.shuffle(labelsets[2])
+
+    #Splits
+    splits = 0.8
+    reliable_train = int(len(labelsets[0])*0.8)
+    unreliable_train = int(len(labelsets[2])*0.8)
+    reliable_val = int(len(labelsets[0])*0.1)
+    unreliable_val = int(len(labelsets[2])*0.1)
+    
+    # WE WILL ADJUST LABEL --> 0 is fake, 1 is true
+    # Like share count is not provided...
+    # Maybe we can extract it by looking up the status ID
+    self.metadata["train"]["crawl"] += [(item, 0,0,1) for item in labelsets[0][:reliable_train]]
+    self.metadata["train"]["crawl"] += [(item, 0,0,0) for item in labelsets[2][:unreliable_train]]
+    random.shuffle(self.metadata["train"]["crawl"])
+
+    self.metadata["val"]["crawl"] += [(item, 0,0,1) for item in labelsets[0][reliable_train:reliable_train+reliable_val]]
+    self.metadata["val"]["crawl"] += [(item, 0,0,0) for item in labelsets[2][unreliable_train:unreliable_train+unreliable_val]]
+    random.shuffle(self.metadata["val"]["crawl"])
+
+    self.metadata["test"]["crawl"] += [(item, 0,0,1) for item in labelsets[0][reliable_train+reliable_val:]]
+    self.metadata["test"]["crawl"] += [(item, 0,0,0) for item in labelsets[2][unreliable_train+unreliable_val:]]
+    random.shuffle(self.metadata["test"]["crawl"])
+      
+
+    self.metadata["train"]["classes"] = self.classes
+    self.metadata["test"]["classes"] = self.classes
+    self.metadata["val"]["classes"] = self.classes
diff --git a/profiles/FNC/experts/crawlers/readme b/profiles/FNC/experts/crawlers/readme
new file mode 100644
index 00000000..10494821
--- /dev/null
+++ b/profiles/FNC/experts/crawlers/readme
@@ -0,0 +1,80 @@
+Each crawler should yield a list of tuples, with the following entries:
+
+
+( full_text, share/retweet, like, label, *data)
+
+We have:
+
+EFND
+Fakeddit
+MiDAS
+
+
+Here, *data are any other labels provided by the 
+dataset. For example, EFND yields a dataset label
+tied to the dataset index. 
+
+
+Generator:
+We will use the HFGenerator. So we need to finish the following options:
+
+No masking
+Token masking (Live generated)
+Word masking (Live generated)
+Keyword masking (pregenerated)
+Keytoken masking (pregenerated)
+
+Verify they all work!!!!
+
+
+----------------------------------------------------------
+Models
+
+
+-----------------------------------------------------------
+Deployments
+
+TF-IDF determiner
+Attention Determiner
+----------------------------------------------------------
+
+Generalization: 
+
+datasets contain complete knowledge
+
+
+datasets contain timeless knowledge
+
+
+B1 -- KO with BERT
+
+July 2021
+
+Introduction --> face recognition and car evolution for KO (people aging)
+
+Artificial Aging
+>>>
+Explicit  <<Charles Isbell>>
+<<<
+
+
+PR solution -- detect bias before it happens
+
+CarZam -->
+
+ADM 
+
+------------------------------------------------
+
+Finding Matching context...?????
+
+
+------------------------------------------------
+
+
+CVPR masking...?
+
+<key is to tell people about the difference between true novelty and stationary changes>
+
+rejuvenation adds new knowledge
+adaptation uses only stationary models
\ No newline at end of file
diff --git a/profiles/FNC/experts/debug/masking_debug.yml b/profiles/FNC/experts/debug/masking_debug.yml
new file mode 100644
index 00000000..84634f55
--- /dev/null
+++ b/profiles/FNC/experts/debug/masking_debug.yml
@@ -0,0 +1,95 @@
+EXECUTION:
+  DATAREADER:
+    crawler_comment: We will use one of the efnd datasets here.
+    CRAWLER_ARGS: 
+      data_folder: Data
+      include: ["kagglefn_long"]
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 500
+      shard_replace: True
+      shardpath: datashard-artifacts
+      shardname: efnd-shard-debug
+      mlm_probability: 0.15
+      masking: True
+      keyword_mask: True
+      keywords: ["doll"]
+      keytoken_mask: False
+      keytokens: ["▁better", "▁blow"]
+      word_mask: False
+      token_mask: False
+      shuffle: False
+      data_shuffle: False
+    GENERATOR: HFGenerator
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  EPOCHS: 3
+  FP16: False
+  MODEL_SERVING: False
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  SKIPEVAL: False
+  TEST_FREQUENCY: 1
+  TRAINER: BaseTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 8
+
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-debug
+  MODEL_BACKBONE: albert
+  MODEL_QUALIFIER: tweets-2020-01-22
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 1 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: HFAutoModel
+  MODEL_BASE: Albert
+  MODEL_KWARGS: 
+    from_pretrained: albert-base-v2
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
+
+
+LOSS: 
+  - LOSSES: ['TorchLoss']
+    KWARGS: 
+      - loss_class: CrossEntropyLoss
+        loss_kwargs:
+          ignore_index: -1
+    LAMBDAS: [1.0]
+    NAME: mask_lm
+    LABEL: mask_lm
+
+OPTIMIZER:
+  - OPTIMIZER: AdamW
+    BASE_LR: 1.0e-5
+    LR_BIAS_FACTOR: 1.0
+    OPTIMIZER_KWARGS: 
+      eps: 1.0e-6
+    OPTIMIZER_NAME: opt-1
+    WEIGHT_BIAS_FACTOR: 0.0005
+    WEIGHT_DECAY: 0.0005
+
+SCHEDULER:
+  - LR_KWARGS: 
+      step_size: 5
+      gamma: 0.5
+    LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
+
diff --git a/profiles/FNC/experts/plugin_generation.py b/profiles/FNC/experts/plugin_generation.py
new file mode 100644
index 00000000..bbbec010
--- /dev/null
+++ b/profiles/FNC/experts/plugin_generation.py
@@ -0,0 +1,16 @@
+from ednaml.deploy import BaseDeploy
+import ednaml.core.decorators as edna
+
+
+@edna.register_deployment
+class FNCPluginDeployment(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    prediction_scores, pooled_out, outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    
+    return prediction_scores, pooled_out, outputs
+  def output_setup(self, **kwargs):
+    pass
+  def output_step(self, logits, features, secondary):
+    pass
\ No newline at end of file
diff --git a/profiles/FNC/extension/config.yml b/profiles/FNC/extension/config.yml
new file mode 100644
index 00000000..2dd347bc
--- /dev/null
+++ b/profiles/FNC/extension/config.yml
@@ -0,0 +1,106 @@
+EXECUTION:
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-filtermask-shard
+      mlm_probability: 0.15
+      masking: False
+      shuffle: True
+      data_shuffle: True
+      keyword_mask: True
+      keywords: ["covid", "corona", "mask", "wuhan", "n95", "sars", "monkey", "pandemic", "social", "quarantin", "virus", "infect", "lock", "ppe", "variant", "vaccine", "travel", "omicron", "ivermectin", 
+        "plandemic", "5g", "gates", "hoax", "bioweapon", "bat", "fauci"]
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  EPOCHS: 3
+  FP16: False
+  MODEL_SERVING: False
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  SKIPEVAL: False
+  TEST_FREQUENCY: 1
+  TRAINER: BaseTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 8
+
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-extension
+  MODEL_BACKBONE: albert
+  MODEL_QUALIFIER: tweets-2020-01-22
+  DRIVE_BACKUP: True
+  SAVE_FREQUENCY: 1 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: FNCAlbertModeler
+  MODEL_BASE: Albert
+  MODEL_KWARGS: 
+    vocab_size_or_config_json_file: 30000
+    embedding_size: 128
+    hidden_size: 768
+    num_hidden_layers: 12
+    num_hidden_groups: 1
+    num_attention_heads: 12
+    intermediate_size: 3072
+    inner_group_num: 1
+    hidden_act: gelu
+    hidden_dropout_prob: 0
+    attention_probs_dropout_prob: 0
+    max_position_embeddings: 512
+    type_vocab_size: 2
+    initializer_range: 0.02
+    layer_norm_eps: 1.0e-12
+    pooling: pooled
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
+
+
+LOSS: 
+  - LOSSES: ['TorchLoss']
+    KWARGS: 
+      - loss_class: CrossEntropyLoss
+        loss_kwargs:
+          ignore_index: -1
+    LAMBDAS: [1.0]
+    NAME: mask_lm
+    LABEL: mask_lm
+
+OPTIMIZER:
+  - OPTIMIZER: AdamW
+    BASE_LR: 1.0e-5
+    LR_BIAS_FACTOR: 1.0
+    OPTIMIZER_KWARGS: 
+      eps: 1.0e-6
+    OPTIMIZER_NAME: opt-1
+    WEIGHT_BIAS_FACTOR: 0.0005
+    WEIGHT_DECAY: 0.0005
+
+SCHEDULER:
+  - LR_KWARGS: 
+      step_size: 5
+      gamma: 0.5
+    LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
+
diff --git a/profiles/FNC/extension/fastkmeans.yml b/profiles/FNC/extension/fastkmeans.yml
new file mode 100644
index 00000000..48f87ba2
--- /dev/null
+++ b/profiles/FNC/extension/fastkmeans.yml
@@ -0,0 +1,42 @@
+
+DEPLOYMENT:
+  DEPLOY: BaseDeploy
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-unmasked-shard
+      mlm_probability: 0.15
+      masking: False
+      shuffle: False
+      data_shuffle: False
+      keyword_mask: False
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: warmup
+    RESET: True
+
+
+MODEL_PLUGIN:
+  - PLUGIN: FastKMP
+    PLUGIN_NAME: FastKMP-l2
+    PLUGIN_KWARGS:
+      proxies: 20
+      dimensions: 768
+      dist: euclidean
+      iterations: 30
+      batch_size: 256
+      alpha: 0.5
\ No newline at end of file
diff --git a/profiles/FNC/extension/neighbors.yml b/profiles/FNC/extension/neighbors.yml
new file mode 100644
index 00000000..ce35f5f7
--- /dev/null
+++ b/profiles/FNC/extension/neighbors.yml
@@ -0,0 +1,45 @@
+
+DEPLOYMENT:
+  OUTPUT_ARGS:
+    neighbor_output: neighbor
+    unfiltered_output: unfiltered
+    basename: tweets-2020-01-22
+  DEPLOY: BaseDeploy
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-unmasked-shard
+      mlm_probability: 0.15
+      masking: False
+      shuffle: False
+      data_shuffle: False
+      keyword_mask: False
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: activated
+    RESET: False
+
+MODEL_PLUGIN:
+  - PLUGIN: FastKMP
+    PLUGIN_NAME: FastKMP-l2
+    PLUGIN_KWARGS:
+      proxies: 20
+      dimensions: 768
+      dist: euclidean
+      iterations: 30
+      batch_size: 256
+      alpha: 0.5
\ No newline at end of file
diff --git a/profiles/FNC/extension/oracle.yml b/profiles/FNC/extension/oracle.yml
new file mode 100644
index 00000000..ec70f4e1
--- /dev/null
+++ b/profiles/FNC/extension/oracle.yml
@@ -0,0 +1,45 @@
+
+DEPLOYMENT:
+  OUTPUT_ARGS:
+    oracle_output: oracle
+    extended_output: extended
+    basename: tweets-2020-01-22
+  DEPLOY: BaseDeploy
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-extended-unmasked-shard
+      mlm_probability: 0.15
+      masking: False
+      shuffle: False
+      data_shuffle: False
+      keyword_mask: False
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: activated
+    RESET: False
+
+MODEL_PLUGIN:
+  - PLUGIN: FastKMP
+    PLUGIN_NAME: FastKMP-l2
+    PLUGIN_KWARGS:
+      proxies: 20
+      dimensions: 768
+      dist: euclidean
+      iterations: 30
+      batch_size: 256
+      alpha: 0.5
\ No newline at end of file
diff --git a/profiles/FNC/extension/training_features.yml b/profiles/FNC/extension/training_features.yml
new file mode 100644
index 00000000..b8de6a75
--- /dev/null
+++ b/profiles/FNC/extension/training_features.yml
@@ -0,0 +1,28 @@
+DEPLOYMENT:
+  OUTPUT_ARGS:
+    feature_file: training_features
+  DEPLOY: BaseDeploy
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-unmasked-shard
+      mlm_probability: 0.15
+      masking: False
+      shuffle: False
+      data_shuffle: False
+      keyword_mask: False
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+  FP16: False
+  EPOCHS: 1
diff --git a/profiles/FNC/filtered/config.yml b/profiles/FNC/filtered/config.yml
new file mode 100644
index 00000000..2bce3264
--- /dev/null
+++ b/profiles/FNC/filtered/config.yml
@@ -0,0 +1,45 @@
+DEPLOYMENT:
+  OUTPUT_ARGS:
+    filtered_output: filtered
+    unfiltered_output: unfiltered
+    basename: tweets-2020-01-22
+  DEPLOY: BaseDeploy
+  DATAREADER:
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-raw
+      azfile: tweets-2020-01-22.json.gz
+    DATAREADER: AlbertReader
+    DATASET_ARGS:
+      classificationclass: ['fn_label']
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-raw-shard
+      shardsize: 20000
+      shuffle: False
+  FP16: False
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-filter
+  MODEL_BACKBONE: offshelf
+  MODEL_QUALIFIER: raw
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 5 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+TRANSFORMATION:
+  BATCH_SIZE: 128
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: FNCFilter
+  MODEL_BASE: NoBase
+  MODEL_KWARGS: 
+    filter_list: ["covid", "corona", "mask", "wuhan", "n95", "sars", "monkey", "pandemic", "social", "quarantin", "virus", "infect", "lock", "ppe", "variant", "vaccine", "travel", "omicron", "ivermectin", 
+      "plandemic", "5g", "gates", "hoax", "bioweapon", "bat", "fauci", ]
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
+
+LOGGING:
+  STEP_VERBOSE: 100 # Batch
\ No newline at end of file
diff --git a/profiles/FNC/fnc-crawler.py b/profiles/FNC/fnc-crawler.py
new file mode 100644
index 00000000..b860a91d
--- /dev/null
+++ b/profiles/FNC/fnc-crawler.py
@@ -0,0 +1,335 @@
+from typing import Any, Dict, List
+import ednaml
+from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
+import json, os
+from ednaml.crawlers import Crawler
+from ednaml.utils.web import download
+from ednaml.utils.file_utils import IterableFile
+import gzip
+import shutil
+
+
+# This is to read from Azure...download a thing from azure, and unzip, and load it into memory...
+# We will datashard it with the Datareader args, specifically for the Deployment...
+
+
+@edna.register_crawler
+class FNCCrawler(Crawler):
+    """Downloads a FNC raw file from Azure if it exists and builds an iterable from the file contents.
+
+    Crawler yields a list of tweet objects, i.e. a dictionary containing tweet attributes.
+    """
+    def __init__(
+        self,
+        logger=None,
+        azstorage="ednadatasets",
+        azcontainer="edna-covid-raw",
+        azfile="tweets-2020-01-22.json.gz",
+        test_file = None,
+        train_azstorage=None,
+        train_azcontainer=None,
+        train_azfile=None,
+        train_file = None
+    ):
+        """Initializes the FNC Crawler.
+
+        By default, only the test data is populated. If the corresponding values for arguments is None, it is not used, instead of throwing error.
+
+        Args:
+            logger (_type_, optional): Logger. Defaults to None.
+
+            azstorage (str, optional): Axure storage name. Defaults to "ednadatasets".
+            azcontainer (str, optional): Azure container name. Defaults to "edna-covid-raw".
+            azfile (str, optional): Azure file name (blob name). Defaults to "tweets-2020-01-22.json.gz".
+            test_file (_type_, optional): If using local files, this will supersede the Azure data. Defaults to None.
+
+            train_azstorage (str, optional): Axure storage name. Defaults to "ednadatasets".
+            train_azcontainer (str, optional): Azure container name. Defaults to "edna-covid-raw".
+            train_azfile (str, optional): Azure file name (blob name). Defaults to "tweets-2020-01-22.json.gz".
+            train_file (_type_, optional): If using local files, this will be present in the "train" metadata and supersede Azure data. Defaults to None.
+            
+        """
+        if test_file is None:
+            if azstorage is not None:
+                az_url = self.build_url(azstorage, azcontainer, azfile)
+                logger.info("Crawling %s" % (az_url))
+                if not os.path.exists(azfile):
+                    download(azfile, az_url)
+                else:
+                    logger.info("%s already exists at %s" % (az_url, azfile))
+                az_jsonfile = os.path.splitext(azfile)[0]
+
+                # Then unzip the file
+                if not os.path.exists(az_jsonfile):
+                    with gzip.open(azfile, "rb") as f_in:
+                        with open(az_jsonfile, "wb") as f_out:
+                            shutil.copyfileobj(f_in, f_out)
+            else:
+                az_jsonfile = None
+        else:   # TODO possible sanity checks????
+            az_jsonfile = test_file
+
+
+        if train_file is None:
+            if train_azstorage is not None:
+                train_az_url = self.build_url(train_azstorage, train_azcontainer, train_azfile)
+                logger.info("Crawling %s" % (train_az_url))
+                if not os.path.exists(train_azfile):
+                    download(train_azfile, train_az_url)
+                else:
+                    logger.info("%s already exists at %s" % (train_az_url, train_azfile))
+                train_az_jsonfile = os.path.splitext(train_azfile)[0]
+
+                # Then unzip the file
+                if not os.path.exists(train_az_jsonfile):
+                    with gzip.open(train_azfile, "rb") as f_in:
+                        with open(train_az_jsonfile, "wb") as f_out:
+                            shutil.copyfileobj(f_in, f_out)
+            else:
+                train_az_jsonfile = None
+        else:   # TODO possible sanity checks????
+            train_az_jsonfile = train_file
+
+        # set up class metadata
+        self.classes = {}
+        self.classes["fnews"] = 2
+
+        # set up content metadata
+        self.metadata = {}
+        self.metadata["secondary"] = {}
+
+        if az_jsonfile is not None:
+            self.metadata["secondary"]["linecount"] = self.bufcount(az_jsonfile)
+        else:
+            self.metadata["secondary"]["linecount"] = 0
+        if train_az_jsonfile is not None:
+            self.metadata["secondary"]["train_linecount"] = self.bufcount(train_az_jsonfile)
+        else:
+            self.metadata["secondary"]["train_linecount"] = 0
+
+        self.metadata["train"] = {}
+        self.metadata["test"] = {}
+        self.metadata["val"] = {}
+        self.metadata["train"]["crawl"] = []
+        self.metadata["test"]["crawl"] = []
+        self.metadata["val"]["crawl"] = []
+
+        # Get the tweet data as list, <id, text, url> TODO
+        # basically need a generator that yields results and maps them through some function
+        # This will be slow because of json.loads line by line. We should speed this up with chunking?
+
+        # TODO extracts full_text only. Fix this so it extracts all relevant objects (or possibly all objects??)
+        if az_jsonfile is not None:
+            self.metadata["test"]["crawl"] = IterableFile(
+                az_jsonfile,
+                #line_callback=lambda row: (json.loads(row.strip())["full_text"], 0),
+                line_callback=lambda row: json.loads(row.strip()),
+            )
+        if train_az_jsonfile is not None:
+            self.metadata["train"]["crawl"] = IterableFile(
+                train_az_jsonfile,
+                #line_callback=lambda row: (json.loads(row.strip())["full_text"], 0),
+                line_callback=lambda row: json.loads(row.strip()),
+            )
+        self.metadata["train"]["classes"] = self.classes
+        self.metadata["test"]["classes"] = self.classes
+        self.metadata["val"]["classes"] = self.classes
+        # pdb.set_trace()
+
+    def build_url(self, azstorage, azcontainer, azfile):
+        return "https://{azstorage}.blob.core.windows.net/{azcontainer}/{azfile}".format(
+            azstorage=azstorage, azcontainer=azcontainer, azfile=azfile
+        )
+
+    def bufcount(self, filename):
+        f = open(filename)                  
+        lines = 0
+        buf_size = 1024 * 1024
+        read_f = f.read # loop optimization
+
+        buf = read_f(buf_size)
+        while buf:
+            lines += buf.count('\n')
+            buf = read_f(buf_size)
+
+        return lines
+
+
+
+import torch
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.generators import TextGenerator
+from torch.utils.data import Dataset as TorchDataset
+import torch.utils.data
+import logging
+from glob import glob
+
+class FNCRawDataset(TorchDataset):
+    """FNCRawDataset converts the raw file into datashards.
+
+    Even if there is no tokenizer performing any conversion, our files are large
+    and may not fit in memory in many cases.
+
+    So, FNCRawDataset will iterate through the Crawler's file iterable, save shards in individual files
+    and record the overall length.
+
+    Since there is no tokenization, there is also no shuffling, since we are not training.
+
+
+    Args:
+        TorchDataset (_type_): _description_
+    """
+    def __init__(self, logger: logging.Logger, crawler_list: IterableFile, mode = "test", **kwargs):
+        self.dataset = crawler_list
+        self.logger = logger
+
+        self.shardsize = kwargs.get("shardsize", 5000)
+        self.shardsaveindex = 0
+        self.shard_replace = kwargs.get("shard_replace", False)
+
+
+        self.shardpath = kwargs.get("shardpath", "datashard-artifacts")
+        self.shardpath = mode+"-"+self.shardpath    #e.g. test-datashard-artifacts
+
+        self.shardname = kwargs.get("shardname", "fnc-shard") + "-"  #the dash
+        self.base_shardpath = os.path.join(self.shardpath, self.shardname)
+        self.shards_exist = False
+
+        if os.path.exists(self.base_shardpath + "0.pt"):
+            if self.shard_replace:
+                self.logger.debug("Deleting existing shards")
+                shutil.rmtree(self.shardpath)
+            else:
+                with open(os.path.join(self.shardpath,"len.txt"), "r") as f:
+                    self.len = int(f.read().strip())
+                self.shards_exist = True
+                self.logger.debug("Shards already exist and `shard_replace` is False")
+        os.makedirs(self.shardpath, exist_ok=True)
+
+        # to track what we have used __getitem__ for so far...
+        self.getcount = 0
+        # Not sure what to do with this yet...
+        self.refresh_flag = self.shardsize
+
+        if not self.shards_exist:
+            self.shardsaveindex = self.sharded_convert_to_features(self.dataset)    # save shards and get numshards
+        else:
+            self.shardsaveindex = len(glob(os.path.join(self.shardpath, "*.pt")))-1   # Get the total number of shards.
+
+        self.shard_load_index = 0   # self.shardsaveindex is the maximum number of shards
+        self.shard_shuffle = list(range(self.shardsaveindex))    # count started from 0. We will not shuffle this
+        #if self.data_shuffle:
+        #    random.shuffle(self.shard_shuffle)
+        self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+        self.current_shardsize = len(self.sharded_dataset)  # because load_shard returns a list...
+        self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0; we will not shuffle this.
+        
+
+
+    def sharded_convert_to_features(self, dataset: IterableFile):
+        """We convert the dataset, an IterableFile, into shards here. Each shard is basically a 
+        json file containing a fixed number of elements.
+
+        Args:
+            dataset (IterableFile): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        shardsaveindex = 0
+        #pbar = tqdm(total=int(len(dataset)/self.shardsize)+1)
+        trow = []
+        for idx, row in enumerate(dataset):
+            trow.append(row)
+
+            if (idx+1)%self.shardsize == 0: # we need to save this shard.
+                
+                self.save_shard(shard=trow,
+                        shard_index = shardsaveindex)
+                shardsaveindex += 1
+                trow = []
+                #pbar.update(1)
+        # final shard...
+        
+        if len(trow) > 0:
+            # SAVE HERE
+            self.save_shard(shard=trow,
+                    shard_index = shardsaveindex)
+            shardsaveindex += 1
+            trow = []
+            #pbar.update(1)
+        #pbar.close()
+        self.len = idx
+        with open(os.path.join(self.shardpath,"len.txt"), "w") as f:
+            f.write(str(self.len))
+        return shardsaveindex
+
+    def save_shard(self, shard: List[Dict[str,str]], shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        with open(shard_save_path, 'w') as f:
+            json.dump(shard, f, ensure_ascii=False)
+        #json.dump(shard, shard_save_path)
+
+    def load_shard(self, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        self.logger.debug("Loading shard %s"%shard_save_path)
+        with open(shard_save_path, 'r') as f:
+          jload = json.load(f)
+        return jload
+
+
+
+
+    def __len__(self):
+        return self.len
+
+    def __getitem__(self, idx):
+        self.getcount += 1
+        if self.getcount == self.current_shardsize:   # we have exhausted examples in this shard
+            self.getcount = 0
+            self.shard_load_index += 1                  # increment the shard index that we will load
+            if self.shard_load_index == self.shardsaveindex: # we have processed all shards.
+                self.shard_load_index = 0
+                #if self.data_shuffle:
+                #    random.shuffle(self.shard_shuffle)
+
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])   #There is no shuffle, so we are fine...
+            self.current_shardsize = len(self.sharded_dataset)
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            #if self.data_shuffle:
+            #    random.shuffle(self.shard_internal_shuffle)
+        #shardindex = idx % self.current_shardsize
+        #return self.sharded_dataset[self.shard_internal_shuffle[shardindex]]
+        return self.sharded_dataset[self.shard_internal_shuffle[self.getcount]]
+
+@edna.register_generator
+class FNCRawGenerator(TextGenerator):
+    """The FNCRawGenerator simply yields the raw text from FNCCrawler in batches. It does no processing through any encoder.
+
+    Args:
+        TextGenerator (_type_): _description_
+    """
+
+    def build_transforms(self, transforms: Dict[str, Any], mode, **kwargs):
+        return None
+
+    def buildDataset(self, datacrawler, mode: str, transform: List[object], **kwargs) -> TorchDataset:
+        return FNCRawDataset(self.logger, datacrawler.metadata[mode]["crawl"], mode, **kwargs)
+
+    def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
+        return torch.utils.data.DataLoader(
+            dataset,batch_size=batch_size*self.gpus,
+                    shuffle=(True if mode=="train" else kwargs.get("shuffle", False)), num_workers = self.workers,
+                    collate_fn=self.collate_fn)
+    
+    def getNumEntities(self, crawler, mode, **kwargs):  #<-- dataset args
+        label_dict = {
+            item: {"classes": crawler.metadata[mode]["classes"][item]}
+            for item in kwargs.get("classificationclass", ["fn_label"])
+        }
+        return LabelMetadata(label_dict=label_dict)
+
+
+    def collate_fn(self, batch):
+        return batch
diff --git a/profiles/FNC/fnc-extended-labeling-crawler.py b/profiles/FNC/fnc-extended-labeling-crawler.py
new file mode 100644
index 00000000..b4697e92
--- /dev/null
+++ b/profiles/FNC/fnc-extended-labeling-crawler.py
@@ -0,0 +1,154 @@
+from typing import Any, Dict, List
+import ednaml
+from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
+import json, os
+from ednaml.crawlers import Crawler
+from ednaml.utils.web import download
+from ednaml.utils.file_utils import IterableFile
+import gzip
+import shutil
+
+
+# This is to read from Azure...download a thing from azure, and unzip, and load it into memory...
+# We will datashard it with the Datareader args, specifically for the Deployment...
+
+
+@edna.register_crawler
+class FNCCrawler(Crawler):
+    """Downloads a FNC raw file from Azure if it exists and builds an iterable from the file contents.
+
+    Crawler yields a list of tweet objects, i.e. a dictionary containing tweet attributes.
+    """
+    def __init__(
+        self,
+        logger=None,
+        azstorage="ednadatasets",
+        azcontainer="edna-covid-raw",
+        azfile="tweets-2020-01-22.json.gz",
+        test_file = None,
+        train_azstorage=None,
+        train_azcontainer=None,
+        train_azfile=None,
+        train_file = None
+    ):
+        """Initializes the FNC Crawler.
+
+        By default, only the test data is populated. If the corresponding values for arguments is None, it is not used, instead of throwing error.
+
+        Args:
+            logger (_type_, optional): Logger. Defaults to None.
+
+            azstorage (str, optional): Axure storage name. Defaults to "ednadatasets".
+            azcontainer (str, optional): Azure container name. Defaults to "edna-covid-raw".
+            azfile (str, optional): Azure file name (blob name). Defaults to "tweets-2020-01-22.json.gz".
+            test_file (_type_, optional): If using local files, this will supersede the Azure data. Defaults to None.
+
+            train_azstorage (str, optional): Axure storage name. Defaults to "ednadatasets".
+            train_azcontainer (str, optional): Azure container name. Defaults to "edna-covid-raw".
+            train_azfile (str, optional): Azure file name (blob name). Defaults to "tweets-2020-01-22.json.gz".
+            train_file (_type_, optional): If using local files, this will be present in the "train" metadata and supersede Azure data. Defaults to None.
+            
+        """
+        if test_file is None:
+            if azstorage is not None:
+                az_url = self.build_url(azstorage, azcontainer, azfile)
+                logger.info("Crawling %s" % (az_url))
+                if not os.path.exists(azfile):
+                    download(azfile, az_url)
+                else:
+                    logger.info("%s already exists at %s" % (az_url, azfile))
+                az_jsonfile = os.path.splitext(azfile)[0]
+
+                # Then unzip the file
+                if not os.path.exists(az_jsonfile):
+                    with gzip.open(azfile, "rb") as f_in:
+                        with open(az_jsonfile, "wb") as f_out:
+                            shutil.copyfileobj(f_in, f_out)
+            else:
+                az_jsonfile = None
+        else:   # TODO possible sanity checks????
+            az_jsonfile = test_file
+
+
+        if train_file is None:
+            if train_azstorage is not None:
+                train_az_url = self.build_url(train_azstorage, train_azcontainer, train_azfile)
+                logger.info("Crawling %s" % (train_az_url))
+                if not os.path.exists(train_azfile):
+                    download(train_azfile, train_az_url)
+                else:
+                    logger.info("%s already exists at %s" % (train_az_url, train_azfile))
+                train_az_jsonfile = os.path.splitext(train_azfile)[0]
+
+                # Then unzip the file
+                if not os.path.exists(train_az_jsonfile):
+                    with gzip.open(train_azfile, "rb") as f_in:
+                        with open(train_az_jsonfile, "wb") as f_out:
+                            shutil.copyfileobj(f_in, f_out)
+            else:
+                train_az_jsonfile = None
+        else:   # TODO possible sanity checks????
+            train_az_jsonfile = train_file
+
+        # set up class metadata
+        self.classes = {}
+        self.classes["fnews"] = 2
+
+        # set up content metadata
+        self.metadata = {}
+        self.metadata["secondary"] = {}
+
+        if az_jsonfile is not None:
+            self.metadata["secondary"]["linecount"] = self.bufcount(az_jsonfile)
+        else:
+            self.metadata["secondary"]["linecount"] = 0
+        if train_az_jsonfile is not None:
+            self.metadata["secondary"]["train_linecount"] = self.bufcount(train_az_jsonfile)
+        else:
+            self.metadata["secondary"]["train_linecount"] = 0
+
+        self.metadata["train"] = {}
+        self.metadata["test"] = {}
+        self.metadata["val"] = {}
+        self.metadata["train"]["crawl"] = []
+        self.metadata["test"]["crawl"] = []
+        self.metadata["val"]["crawl"] = []
+
+        # Get the tweet data as list, <id, text, url> TODO
+        # basically need a generator that yields results and maps them through some function
+        # This will be slow because of json.loads line by line. We should speed this up with chunking?
+
+        # TODO extracts full_text only. Fix this so it extracts all relevant objects (or possibly all objects??)
+        if az_jsonfile is not None:
+            self.metadata["test"]["crawl"] = IterableFile(
+                az_jsonfile,
+                line_callback=lambda row: (json.loads(row.strip())["full_text"], 0,0,0),
+            )
+        if train_az_jsonfile is not None:
+            self.metadata["train"]["crawl"] = IterableFile(
+                train_az_jsonfile,
+                line_callback=lambda row: (json.loads(row.strip())["full_text"], 0,0,0),
+            )
+        self.metadata["train"]["classes"] = self.classes
+        self.metadata["test"]["classes"] = self.classes
+        self.metadata["val"]["classes"] = self.classes
+        # pdb.set_trace()
+
+    def build_url(self, azstorage, azcontainer, azfile):
+        return "https://{azstorage}.blob.core.windows.net/{azcontainer}/{azfile}".format(
+            azstorage=azstorage, azcontainer=azcontainer, azfile=azfile
+        )
+
+    def bufcount(self, filename):
+        f = open(filename)                  
+        lines = 0
+        buf_size = 1024 * 1024
+        read_f = f.read # loop optimization
+
+        buf = read_f(buf_size)
+        while buf:
+            lines += buf.count('\n')
+            buf = read_f(buf_size)
+
+        return lines
\ No newline at end of file
diff --git a/profiles/FNC/fnc-extension.py b/profiles/FNC/fnc-extension.py
new file mode 100644
index 00000000..a29b1f23
--- /dev/null
+++ b/profiles/FNC/fnc-extension.py
@@ -0,0 +1,403 @@
+"""
+Albert model with a MLM head
+
+"""
+
+import ednaml
+from ednaml.crawlers import Crawler
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.models.Albert import AlbertPreTrainedModel, AlbertModel, AlbertOnlyMLMHead
+from ednaml.models.Albert import AlbertConfig
+from torch.nn import CrossEntropyLoss
+from ednaml.models.ModelAbstract import ModelAbstract
+import torch
+from torch import nn
+import ednaml.core.decorators as edna
+
+from ednaml.utils.layers import GradientReversalLayer
+from ednaml.utils.albert import (
+    AlbertEmbeddingAverage,
+    AlbertPooledOutput,
+    AlbertRawCLSOutput,
+)
+
+
+@edna.register_model
+class FNCAlbertModeler(ModelAbstract):
+    def __init__(
+        self, base, weights, metadata, normalization, parameter_groups, **kwargs
+    ):
+        super().__init__(
+            base=base,
+            weights=weights,
+            metadata=metadata,
+            normalization=normalization,
+            parameter_groups=parameter_groups,
+            **kwargs
+        )
+
+    def model_attributes_setup(self, **kwargs):
+        self.config = AlbertConfig(**kwargs)
+        self.configargs = self.config.getVars()
+        self.pool_method = kwargs.get("pooling", "pooled")
+
+    def model_setup(self, **kwargs):
+        self.encoder, errors = AlbertModel.from_pretrained(
+            "pytorch_model.bin", config=self.config, output_loading_info=True
+        )
+        print("Errors \n\t", errors)
+        self.decoders = AlbertOnlyMLMHead(self.config)
+        self.decoders.apply(self._init_weights)
+
+        self.tie_weights()
+
+        if self.pool_method == "pooled":
+            self.pooler_layer = AlbertPooledOutput()
+        elif self.pool_method == "raw":
+            self.pooler_layer = AlbertRawCLSOutput()
+        elif self.pool_method == "average":
+            self.pooler_layer = AlbertEmbeddingAverage()
+        else:
+            raise NotImplementedError()
+
+    def tie_weights(self):
+        self._tie_or_clone_weights(
+            self.decoders, self.encoder.embeddings.word_embeddings
+        )
+
+    def _tie_or_clone_weights(self, first_module, second_module):
+        """ Tie or clone module weights depending of weither we are using TorchScript or not
+        """
+
+        if self.config.torchscript:
+            first_module.weight = nn.Parameter(second_module.weight.clone())
+        else:
+            first_module.weight = second_module.weight
+
+        if hasattr(first_module, "bias") and first_module.bias is not None:
+            first_module.bias.data = torch.nn.functional.pad(
+                first_module.bias.data,
+                (0, first_module.weight.shape[0] - first_module.bias.shape[0]),
+                "constant",
+                0,
+            )
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, torch.nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+    def forward_impl(
+        self,
+        x,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+    ):
+        outputs = self.encoder(
+            x,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+        )  # sequence_output, pooled_output, (hidden_states), (attentions)
+
+        sequence_output = outputs[0]
+        prediction_scores = self.decoders(sequence_output)
+        if self.inferencing:
+            return prediction_scores, self.pooler_layer(outputs)
+        pooled_out = self.pooler_layer(outputs)
+
+        return (
+            prediction_scores,
+            pooled_out,
+            outputs,
+        )  # list of k scores; hidden states, attentions...
+
+    def partial_load(self, weights_path):
+        super().partial_load(
+            self, weights_path
+        )  # For this, we need to look at the from_pretrained function to accurately load the saved weights from .bin...
+
+
+from ednaml.trainer import BaseTrainer
+from torch.utils.data import DataLoader
+from logging import Logger
+from typing import Dict, List
+from ednaml.config.EdnaMLConfig import EdnaMLConfig
+import tqdm
+
+
+@edna.register_trainer
+class FNCAlbertModelerTrainer(BaseTrainer):
+    def __init__(
+        self,
+        model: ModelAbstract,
+        loss_fn: List[ednaml.loss.builders.LossBuilder],
+        optimizer: List[torch.optim.Optimizer],
+        loss_optimizer: List[torch.optim.Optimizer],
+        scheduler: List[torch.optim.lr_scheduler._LRScheduler],
+        loss_scheduler: List[torch.optim.lr_scheduler._LRScheduler],
+        train_loader: DataLoader,
+        test_loader: DataLoader,
+        epochs: int,
+        skipeval: bool,
+        logger: Logger,
+        crawler: Crawler,
+        config: EdnaMLConfig,
+        labels: LabelMetadata,
+        **kwargs
+    ):
+        super().__init__(
+            model,
+            loss_fn,
+            optimizer,
+            loss_optimizer,
+            scheduler,
+            loss_scheduler,
+            train_loader,
+            test_loader,
+            epochs,
+            skipeval,
+            logger,
+            crawler,
+            config,
+            labels,
+            **kwargs
+        )
+        self.maskedaccuracy = []
+        self.fullaccuracy = []
+
+    def step(self, batch):
+        batch = tuple(item.cuda() for item in batch)
+        (
+            all_input_ids,
+            all_attention_mask,
+            all_token_type_ids,
+            all_masklm,
+            all_labels,
+        ) = batch
+        outputs = self.model(
+            all_input_ids,
+            token_type_ids=all_token_type_ids,
+            attention_mask=all_attention_mask,
+        )
+
+        pred = outputs[0]
+        discrimination = outputs[1]
+        l_d_logits = outputs[2]
+
+        # get the size of predictions, then sum masked_loss
+        # We will late deal with doing the masked_lm for specific decoders only
+        # That, is, compute cross-entropy, then zero out the errors when the dataset does not match decoder
+        masked_loss = None
+        masked_pred = None 
+        masked_acc = None
+        
+        masked_loss = self.loss_fn["mask_lm"](
+            input=pred.view(
+                -1, self.model.config.vocab_size
+            ),
+            target=all_masklm.view(-1),
+        )
+        masked_pred = (
+            (
+                pred
+                .view(-1, self.model.config.vocab_size)
+                .max(1)[1]
+                == all_input_ids.view(-1)
+            )
+            .float()
+            .mean()
+            .cpu()
+        )
+        masked_idx = all_masklm.view(-1) == -1
+        masked_acc = (
+            (
+                pred
+                .view(-1, self.model.config.vocab_size)[masked_idx]
+                .max(1)[1]
+                == all_input_ids.view(-1)[masked_idx]
+            )
+            .float()
+            .mean()
+            .cpu()
+        )
+        # zero out the bad losses here...
+        lm_loss = masked_loss
+        
+    
+
+        self.losses["mask_lm"].append(lm_loss.item())
+        masked_pred = masked_pred if not torch.isnan(masked_pred) else 0 
+        masked_acc = masked_acc if not torch.isnan(masked_acc) else 0 
+        self.fullaccuracy.append(masked_pred)
+        self.maskedaccuracy.append(masked_acc)
+
+        return lm_loss
+
+    def evaluate_impl(self):
+        # 1 because single domain
+        
+        prediction_acc, maskedlm_acc, unmaskedlm_acc = [], [], []
+        with torch.no_grad():
+            for batch in tqdm.tqdm(
+                self.test_loader, total=len(self.test_loader), leave=False
+            ):
+                batch = tuple(item.cuda() for item in batch)
+                (
+                    all_input_ids,
+                    all_attention_mask,
+                    all_token_type_ids,
+                    all_masklm,
+                    all_labels
+                ) = batch
+                outputs = self.model(
+                    all_input_ids,
+                    token_type_ids=all_token_type_ids,
+                    attention_mask=all_attention_mask,
+                )
+                pred = outputs[0]
+                
+
+                masked_pred = None
+                masked_acc = None
+                unmasked_acc = None
+                masked_pred = (
+                    (
+                        pred
+                        .view(-1, self.model.config.vocab_size)
+                        .max(1)[1]
+                        == all_input_ids.view(-1)
+                    )
+                    .float()
+                    .mean()
+                    .cpu()
+                )
+                masked_idx = all_masklm.view(-1) == -1
+                unmasked_idx = all_masklm.view(-1) != -1
+                masked_acc = (
+                    (
+                        pred
+                        .view(-1, self.model.config.vocab_size)[masked_idx]
+                        .max(1)[1]
+                        == all_input_ids.view(-1)[masked_idx]
+                    )
+                    .float()
+                    .mean()
+                    .cpu()
+                )
+                unmasked_acc = (
+                    (
+                        pred
+                        .view(-1, self.model.config.vocab_size)[unmasked_idx]
+                        .max(1)[1]
+                        == all_input_ids.view(-1)[
+                            unmasked_idx
+                        ]
+                    )
+                    .float()
+                    .mean()
+                    .cpu()
+                )
+
+
+                masked_pred =  masked_pred if not torch.isnan(masked_pred) else 0 
+                masked_acc =  masked_acc if not torch.isnan(masked_acc) else 0 
+                unmasked_acc =  unmasked_acc if not torch.isnan(unmasked_acc) else 0 
+
+                prediction_acc.append(masked_pred)
+                maskedlm_acc.append(masked_acc)
+                unmaskedlm_acc.append(unmasked_acc)
+
+        p_accs = [torch.mean(torch.Tensor(item)).item() for item in [prediction_acc]]
+        m_accs = [torch.mean(torch.Tensor(item)).item() for item in [maskedlm_acc]]
+        u_accs = [torch.mean(torch.Tensor(item)).item() for item in [unmaskedlm_acc]]
+
+        p_str = "\t\tReconstruction\t" + "\t".join(
+            ["Domain {0}: {1:.3f}".format(idx, item) for idx, item in enumerate(p_accs)]
+        )
+        m_str = "\tMasked Prediction\t" + "\t".join(
+            ["Domain {0}: {1:.3f}".format(idx, item) for idx, item in enumerate(m_accs)]
+        )
+        u_str = "\tUnmasked Prediction\t" + "\t".join(
+            ["Domain {0}: {1:.3f}".format(idx, item) for idx, item in enumerate(u_accs)]
+        )
+
+        self.logger.info(p_str)
+        self.logger.info(m_str)
+        self.logger.info(u_str)
+
+        return None, None, None # maybe return predictions or something ...?
+
+    def printStepInformation(self):
+        loss_avg = [0.0] * len(self.losses)
+        for idx, lossname in enumerate(self.losses):
+            loss_avg[idx] += (
+                sum(self.losses[lossname][-self.step_verbose :]) / self.step_verbose
+            )
+        # loss_avg /= self.num_losses
+        soft_avg = sum(self.fullaccuracy[-100:]) / float(len(self.fullaccuracy[-100:]))
+        mask_avg = sum(self.maskedaccuracy[-100:]) / float(
+            len(self.maskedaccuracy[-100:])
+        )
+        self.logger.info(
+            "Epoch{0}.{1}\tMaskedLM: {2:.3f}\tReconstruct: {3:.3f}\tMasked Acc: {4:.3f}".format(
+                self.global_epoch,
+                self.global_batch,
+                loss_avg[0],
+                soft_avg,
+                mask_avg,
+            )
+        )
+
+
+import ednaml, torch, os, csv
+from ednaml.deploy.BaseDeploy import BaseDeploy
+import h5py
+
+@edna.register_deployment
+class FNCTrainingFeaturesDeploy(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    prediction_scores, pooled_out, outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    
+    return None, pooled_out.cpu(), None
+
+  def output_setup(self, **kwargs):
+    output_file = kwargs.get("feature_file", "training_features")
+    self.output_file = output_file + ".h5"
+    self.writer = h5py.File(self.output_file, "w")  #we will delete old training features file
+    self.written = False
+    self.prev_idx = -1
+    
+  def output_step(self, logits, features: torch.LongTensor, secondary): 
+    if self.written:
+        feats = features.numpy()
+        self.writer["features"].resize((self.writer["features"].shape[0] + feats.shape[0]), axis=0)
+        self.writer["features"][-feats.shape[0]:] = feats        
+    else:   # First time writing -- we will need to create the dataset.
+        self.writer.create_dataset("features", data=features.numpy(), compression = "gzip", chunks=True, maxshape=(None,features.shape[1]))
+        self.written = True
+
+    if self.writer["features"].shape[0]%5000 > self.prev_idx:
+        self.logger.debug("Chunked %i lines in deployment output %s"%(self.writer["features"].shape[0], self.output_file))
+        self.prev_idx = self.writer["features"].shape[0]%5000
+
+  def end_of_epoch(self, epoch: int):
+      self.writer.close()
+
+  def end_of_deployment(self):
+      pass
+
+
+
diff --git a/profiles/FNC/fnc-fastkmeans.py b/profiles/FNC/fnc-fastkmeans.py
new file mode 100644
index 00000000..3cfb2bdd
--- /dev/null
+++ b/profiles/FNC/fnc-fastkmeans.py
@@ -0,0 +1,178 @@
+import ednaml.core.decorators as edna
+
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+from sklearn.cluster import MiniBatchKMeans
+import h5py, os
+from ednaml.plugins.ModelPlugin import ModelPlugin
+
+@edna.register_model_plugin
+class FastKMP(ModelPlugin):
+    name = "FastKMP"
+    def __init__(self, proxies = 10, dimensions = 768, dist = "euclidean", rand_seed = 34623498, iterations = 25, batch_size = 256, alpha = 0.5, feature_file = None):
+        super().__init__(proxies = proxies, dimensions = dimensions, dist = dist, rand_seed = rand_seed, iterations = iterations, batch_size = batch_size, alpha = alpha, feature_file = feature_file)
+
+
+    def build_plugin(self, **kwargs):
+        self.proxies = kwargs.get("proxies")
+        self.rand_seed = kwargs.get("rand_seed")
+        self.dimensions = kwargs.get("dimensions")
+        self.dist = kwargs.get("dist")    # TODO
+        self.batch_size = kwargs.get("batch_size")
+        self.iterations = kwargs.get("iterations")
+        self.alpha = kwargs.get("alpha")
+
+        self.activated = False
+        self._dist_func = self.build_dist(self.dist)
+        if self.dist == "euclidean":
+            self._preprocess = lambda x:x
+        elif self.dist == "cosine":
+            self._preprocess = torch.nn.functional.normalize
+
+        output_file = kwargs.get("feature_file")
+        if output_file is None:
+            import uuid
+            output_file = "-".join(["FKMP", "features", uuid.uuid1().urn.split("-")[0].split(":")[-1]])
+        self._output_file = output_file + ".h5"
+        if os.path.exists(self._output_file):
+            os.remove(self._output_file)
+        self._writer = h5py.File(self._output_file, "w")  #we will delete old training features file
+        self._written = False
+        self._prev_idx = -1
+
+
+    def build_dist(self, dist="euclidean"):
+        if dist == "euclidean":
+            return self.l2dist
+        elif dist == "cosine":
+            return self.cosdist
+        else:
+            raise ValueError("Invalid value for dist: %s"%dist)
+    def l2dist(self, x):
+        return torch.argmin(torch.sqrt(((self.cluster_means - x)**2).sum(1)))
+    def cosdist(self, x):   #https://stackoverflow.com/questions/46409846/using-k-means-with-cosine-similarity-python
+        return self.l2dist(torch.nn.functional.normalize(x, dim=0))
+
+    def build_params(self, **kwargs):
+        self.cluster_means = torch.rand(self.proxies, self.dimensions)
+        self.inertia = 0
+        self.cluster_counts = torch.zeros(self.proxies)
+        self.kdcluster = None
+        self.high_density_thresholds = []
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self.activated:
+            # perform the training here
+            self.save_features(features)
+            
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            dist, labels, idx = self.compute_labels(features)
+            return feature_logits, features, secondary_outputs, kwargs, {"threshold": labels, "distance": dist, "label": idx}
+
+    def save_features(self, features):
+        feats = features.cpu().numpy()
+        if self._written:
+            self._writer["features"].resize((self._writer["features"].shape[0] + feats.shape[0]), axis=0)
+            self._writer["features"][-feats.shape[0]:] = feats        
+        else:   # First time writing -- we will need to create the dataset.
+            self._writer.create_dataset("features", data=feats, compression = "gzip", chunks=True, maxshape=(None,feats.shape[1]))
+            self._written = True
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        try:    # in case it is closed already somehow...
+            self._writer.close()
+        except:
+            pass
+        if not self.activated:  # If already activated, we do not need to change anything
+            self.post_epoch_num = epoch    # Maybe have a better check in case something drops or epochs get skipped...
+            if self.pre_epoch_flag: # Means we had a pre_epoch flag so an epoch was completed.
+                self.pre_epoch_flag = False # reset the flag
+                if self.post_epoch_num != self.pre_epoch_num:
+                    raise ValueError("Epoch may have been skipped. Before: %i\tAfter: %i"%(self.pre_epoch_num, self.post_epoch_num))
+            
+            self.performkmeans()
+            self.highdensitybins(model)
+            
+            self.activated = True
+
+    def performkmeans(self):
+        import time
+        inertia = []
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        
+        print("Starting KMeans for k={kval}, with {iters} iterations".format(kval=self.proxies, iters=self.iterations))
+        kmeans = MiniBatchKMeans(n_clusters = self.proxies, random_state = self.rand_seed, batch_size = self.batch_size)
+        stime = time.time()
+        for iters in range(self.iterations):
+            for i in range(0, data_size, self.batch_size):
+                current_data = data['features'][i:i+self.batch_size]
+                kmeans.partial_fit(current_data)
+            if iters%5 == 0:
+                etime = round(time.time() - stime, 2)
+                print("\t[{elapse} s] -- Completed {iters} iterations".format(iters=iters, elapse = etime))
+                stime = time.time()
+        print("\tCompeted MBKM for k={kval}, with inertia: {inertia}".format(kval=self.proxies, inertia = kmeans.inertia_))
+        data.close()
+        self.cluster_means = torch.tensor(kmeans.cluster_centers_)
+        self.inertia = kmeans.inertia_
+        return inertia
+
+    def highdensitybins(self, model: ModelAbstract):
+        if self.dist == "euclidean":
+            self.kdcluster = KDTree(self.cluster_means)
+        elif self.dist == "cosine":
+            self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+        else:
+            raise ValueError("Invalid value for dist: %s"%self.dist)
+
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        distance_bins = [[] for _ in range(self.proxies)]
+        print("Starting High Density estimation")
+        for i in range(0, data_size, self.batch_size):
+            current_data = data['features'][i:i+self.batch_size]
+            dist, indices = self.kdcluster.query(self._preprocess(torch.tensor(current_data)), k=1, return_distance=True)   #.squeeze()
+            for idx in indices[:,0]:
+                distance_bins[idx].append(dist[idx, 0])
+        
+        self.high_density_thresholds = [None]*self.proxies
+        import numpy as np
+        for proxy in range(self.proxies):
+            self.high_density_thresholds[proxy] = np.percentile(distance_bins[proxy], self.alpha * 100)
+        print("Completed High Density threshold estimation")
+        data.close()
+
+    def compute_labels(self, features):
+        """Compute the cluster labels for dataset X given centers C.
+        """
+        # labels = np.argmin(pairwise_distances(C, X), axis=0) # THIS REQUIRES TOO MUCH MEMORY FOR LARGE X
+        feats = features.cpu()
+        dist, idx = self.kdcluster.query(self._preprocess(feats), k=1, return_distance=True)   #.squeeze()
+        # TODO convert idx to the actual cluster means to the actual cluster labels...
+        return torch.from_numpy(dist).squeeze(1), torch.tensor([self.high_density_thresholds[item[0]] for item in idx]), idx[:,0]
+
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
+
+from ednaml.deploy import BaseDeploy
+
+@edna.register_deployment
+class FNCPluginDeployment(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    prediction_scores, pooled_out, outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    
+    return prediction_scores, pooled_out, outputs
+  def output_setup(self, **kwargs):
+    pass
+  def output_step(self, logits, features, secondary):
+    pass
\ No newline at end of file
diff --git a/profiles/FNC/fnc-filtered.py b/profiles/FNC/fnc-filtered.py
new file mode 100644
index 00000000..0ea2576f
--- /dev/null
+++ b/profiles/FNC/fnc-filtered.py
@@ -0,0 +1,74 @@
+import os, json
+import torch
+from ednaml.deploy.BaseDeploy import BaseDeploy
+from ednaml.models.Albert import AlbertPreTrainedModel, AlbertModel, AlbertOnlyMLMHead
+from ednaml.models.Albert import AlbertConfig
+from torch.nn import CrossEntropyLoss
+from ednaml.models.ModelAbstract import ModelAbstract
+from torch import nn
+
+from ednaml.utils.albert import AlbertEmbeddingAverage, AlbertPooledOutput, AlbertRawCLSOutput
+import ednaml.core.decorators  as edna
+
+
+
+@edna.register_model
+class FNCFilter(ModelAbstract):
+    model_name = "FNCFilter"
+    """FNCFilter takes in the raw FNC data and keeps text documents that contain 
+    specific keywords.
+
+    Documents with these keywords are saved in fnc-filtered.json, and documents without these keywords are saved in fnc-unfiltered.json
+
+    Args:
+        ModelAbstract (_type_): _description_
+    """
+
+    def model_attributes_setup(self, **kwargs):
+        self.filter_list = kwargs.get("filter_list", ["covid", "corona"])
+
+    def model_setup(self, **kwargs):
+        self.filter_func = lambda x : 1 if len([item for item in self.filter_list if item in x]) > 0 else 0
+
+    def forward_impl(self, x, **kwargs): # x is list of objects in batch
+        # we are provided a batch of objects as a list
+        filter_batch = []
+        unfilter_batch = []
+        for item in x:
+            if self.filter_func(item["full_text"].lower()):
+                filter_batch.append(item)
+            else:
+                unfilter_batch.append(item)
+
+        return filter_batch, [], unfilter_batch
+
+        
+
+@edna.register_deployment
+class FNCFilterDeployment(BaseDeploy):
+    def deploy_step(self, batch):   # batch should be list of dicts
+        filter_batch, _, unfilter_batch = self.model(batch)
+
+        return filter_batch, None, unfilter_batch
+
+    def output_setup(self, **kwargs):
+        filter_file = kwargs.get("filtered_output")
+        unfilter_file = kwargs.get("unfiltered_output")
+        basename = kwargs.get("basename")
+        filter_name = "-".join([basename, filter_file])+".json"
+        unfilter_name = "-".join([basename, unfilter_file])+".json"
+
+        self.filter_obj = open(filter_name, "w")
+        self.unfilter_obj = open(unfilter_name, "w")
+
+    def output_step(self, logits, features, secondary):
+        #logits is filter_batch
+        # secondary is unfilter_batch
+        for item in logits:
+            self.filter_obj.write(json.dumps(item)+"\n")
+        for item in secondary:
+            self.unfilter_obj.write(json.dumps(item)+"\n")
+
+    def end_of_deployment(self):
+        self.filter_obj.close()
+        self.unfilter_obj.close()
diff --git a/profiles/FNC/fnc-filtermasked.py b/profiles/FNC/fnc-filtermasked.py
new file mode 100644
index 00000000..f2966175
--- /dev/null
+++ b/profiles/FNC/fnc-filtermasked.py
@@ -0,0 +1,376 @@
+from glob import glob
+import random
+import torch, os, shutil
+from torch.utils.data import TensorDataset
+from tqdm import tqdm
+import ednaml.core.decorators as edna
+
+
+
+
+class FNCFilterMaskDataset(torch.utils.data.Dataset):
+    def __init__(self, logger, dataset, mode, transform=None, **kwargs):
+        self.dataset = dataset  # list of tuples (text, labels, labels)
+        self.logger = logger
+        self.file_len = len(dataset)    #kwargs.get("crawler_secondary").get("linecount", 1)
+        self.data_shuffle = kwargs.get("data_shuffle", True)
+
+        # Options
+        self.cache = kwargs.get("cache", False)
+        self.memcache = kwargs.get("memcache", False)
+
+        self.shardcache = kwargs.get("shardcache", False)
+        self.shardsize = kwargs.get("shardsize", 0)
+        self.shardsaveindex = 0
+        self.shard_replace = kwargs.get("shard_replace", False)
+        self.shardpath = kwargs.get("shardpath", "datashard-artifacts")
+        self.shardpath = mode+"-"+self.shardpath
+        self.shardname = kwargs.get("shardname", "fnc-filtermask-shard") + "-"  #the dash
+        if self.shardcache:
+            self.logger.debug("[Mode `{mode}`] Will look in path [{path}] for shards `{shards}[e].pt`".format(path=self.shardpath, shards=self.shardname, mode=mode))
+        self.base_shardpath = os.path.join(self.shardpath, self.shardname)
+        self.shards_exist = False
+        if os.path.exists(self.base_shardpath + "0.pt"):
+            if self.shard_replace:
+                self.logger.debug("Deleting existing shards")
+                shutil.rmtree(self.shardpath)
+            else:
+                self.shards_exist = True
+                self.logger.debug("Shards already exist and `shard_replace` is False")
+        else:
+          self.logger.debug("Shards do not exist and will be created.")
+        if self.shardcache:
+            self.logger.debug("Creating shardpath %s"%self.shardpath)
+            os.makedirs(self.shardpath, exist_ok=True)
+        
+        
+        
+        self.word_masking = kwargs.get("word_mask", False)
+        self.keyword_masking = kwargs.get("keyword_mask", True)
+        self.token_masking = kwargs.get("token_mask", False)
+        self.keytoken_masking = kwargs.get("keytoken_mask", False)
+        
+        
+        # Tokenizer and lengths
+        self.tokenizer = kwargs.get("tokenizer")
+        self.maxlen = kwargs.get("maxlen")
+        self.mlm = kwargs.get("mlm_probability", 0.2)
+        self.masking = kwargs.get("masking", mode=="train")
+
+
+        self.keywords = []
+        if self.keyword_masking:
+            self.keywords = kwargs.get("keywords", [])
+        self.keytokens = []
+        if self.keytoken_masking:
+            self.keytokens = kwargs.get("keytokens", [])
+        # Store indices corresponding to keytokens
+        vocab = self.tokenizer.get_vocab()
+        self.keytokens = {vocab[item]:1 for item in self.keytokens if item in vocab}
+
+        # This is for memcache and shardcache
+        self.getcount = 0
+        self.shardcount = 0
+        if self.memcache:
+            self.refresh_flag = len(dataset)    # Checks when we have processed entire dataset, refreshes mask ids when needed
+        if self.shardcache:
+            self.refresh_flag = self.shardsize
+            # this shuffles the internal shard index, so that we get a random example, instead of in order
+            self.shard_internal_shuffle = []
+            # This shuffles the shard index, so that each epoch, we load shards in different orders. We will set this up after setting up shards...
+            self.shard_shuffle = []
+
+        if int(self.cache) + int(self.memcache) + int(self.shardcache) > 1:
+            raise ValueError("Use only one cache type")
+        
+        self.getter = self.uncachedget  # for each entry, directly convert to features and return
+        if self.cache:
+            raise NotImplementedError()
+            self.getter = self.diskget
+        elif self.memcache:
+            self.getter = self.memget
+        elif self.shardcache:
+            self.getter = self.shardget
+
+        # if cache, then we will cache transformations to disk, then load them
+        # if memcache, then we will cache to memory and load them
+        # if shardcache, we will cache to disk in a sharded fashion, and load shards
+
+        if self.cache or self.memcache:
+        # The actual cache-ing
+            self.logger.debug("Started mem caching")
+            self.input_length_cache = []
+            self.logger.debug("Converting to features")
+            self.convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen)
+            self.logger.debug("Masking")
+            self.memcached_dataset = self.refresh_mask_ids()
+
+        if self.shardcache:
+            self.input_length_cache = []
+            if not self.shards_exist:
+                self.logger.debug("Generating shards")
+                self.shardsaveindex = self.sharded_convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen)    # save shards and get numshards
+            else:
+                self.shardsaveindex = len(glob(os.path.join(self.shardpath, "*.pt")))-1   # TODO Bug fix if files do not have consistent numbering
+            self.logger.debug("Obtained %i shards"%self.shardsaveindex)
+            self.shard_load_index = 0   # self.shardsaveindex is the maximum number of shards
+            self.shard_shuffle = list(range(self.shardsaveindex))    # count started from 0
+            if self.data_shuffle:
+                self.logger.debug("Shuffling shard load order")
+                random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+            if self.masking:
+                self.logger.debug("Refreshing token masks for loaded shard")    # TODO
+                self.sharded_dataset = self.refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+            self.current_shardsize = len(self.sharded_dataset)
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+
+
+
+    def shardget(self, idx):
+        # get entry from already loaded shardcache. so idx is incremented one by one -- No shuffling in data loader.
+        # This is a design choice, and we can't really do anything if user does shuffle + sharding
+        self.getcount += 1
+        if self.getcount == self.current_shardsize:   # we have exhausted examples in this shard
+            self.getcount = 0
+            self.shard_load_index += 1                  # increment the shard index that we will load
+            if self.shard_load_index == self.shardsaveindex: # we have processed all shards.
+                self.shard_load_index = 0
+                if self.data_shuffle:
+                    random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+
+            if self.masking:
+                self.sharded_dataset = self.refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+
+            self.current_shardsize = len(self.sharded_dataset)
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+        #shardindex = idx % self.current_shardsize
+        return self.sharded_dataset[self.shard_internal_shuffle[self.getcount]]
+
+    def sharded_convert_to_features(self, dataset, tokenizer, maxlen):
+        features = []
+        shardsaveindex = 0
+        self.input_length_cache = []
+        self.logger.info("Generating shards")
+        pbar = tqdm(total=int(self.file_len/self.shardsize)+1)
+        for idx, sample in enumerate(dataset):
+            if self.keyword_masking:
+                word_tokens = sample["full_text"].split(" ")
+                keyword_idx = []
+                for widx,item in enumerate(word_tokens):
+                    if sum([1 if kword in item else 0 for kword in self.keywords]) > 0:
+                        keyword_idx.append(widx) # i.e. find word index for each keyword, if it exists
+
+            encoded = self.tokenizer(sample["full_text"], return_tensors="pt", padding="max_length", max_length = maxlen, truncation = True, return_length = True)
+            enc_len = torch.sum(encoded["attention_mask"])
+            self.input_length_cache.append(enc_len)
+            # create a list of lists. Each sublist is a mask for each keyword. So, we can AND all sublists to get the overall mask.
+            # Edge cases -- no keywords: 
+            if self.keyword_masking:
+                match_idxs = torch.LongTensor([[wid != keyword_idx[kidx] for wid in encoded.word_ids(0)] for kidx in range(len(keyword_idx))])
+                if match_idxs.shape[0] > 0: #i.e. we have a mask for keywords, so we will and everything
+                    match_idxs = torch.all(match_idxs, dim=0, keepdim=True)
+                    merged_mask = torch.where(match_idxs == 0, match_idxs, encoded["attention_mask"])
+                else:
+                    merged_mask = encoded["attention_mask"]
+            else:
+                merged_mask = encoded["attention_mask"]
+            #if self.keytoken_masking:
+
+            
+            features.append(
+                (encoded["input_ids"], merged_mask, encoded["token_type_ids"], enc_len, 0)  # 0 used to be *sample[1:], i.e. the label
+            )
+
+            if (idx+1)%self.shardsize == 0: # we need to save this shard.
+                # TODO need to handle multiple labels, later...
+                all_input_ids = torch.cat([f[0] for f in features])
+                all_attention_mask = torch.cat([f[1] for f in features])
+                all_token_type_ids = torch.cat([f[2] for f in features])
+                all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+                all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+                all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+
+                # Propagate input masking to output masking
+                for midx in range(all_attention_mask.shape[0]):
+                    masked_index = torch.where(all_attention_mask[midx][:self.input_length_cache[midx]]==0)[0]
+                    all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+                self.input_length_cache = []
+
+                # SAVE HERE
+                self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_labels),
+                        shard_index = shardsaveindex)
+                shardsaveindex += 1
+                features = []
+                pbar.update(1)
+        # final shard...
+        if len(features) > 0:
+            all_input_ids = torch.cat([f[0] for f in features])
+            all_attention_mask = torch.cat([f[1] for f in features])
+            all_token_type_ids = torch.cat([f[2] for f in features])
+            all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+            all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+            all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+
+            # Propagate input masking to output masking
+            for midx in range(all_attention_mask.shape[0]):
+                masked_index = torch.where(all_attention_mask[midx][:self.input_length_cache[midx]]==0)[0]
+                all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+            self.input_length_cache = []
+
+            # SAVE HERE
+            self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_labels),
+                    shard_index = shardsaveindex)
+            shardsaveindex += 1
+            features = []
+            pbar.update(1)
+        pbar.close()
+        return shardsaveindex
+
+    def save_shard(self, shard: TensorDataset, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        torch.save(shard, shard_save_path)
+
+    def load_shard(self, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        self.logger.debug("Loading shard %s"%shard_save_path)
+        return torch.load(shard_save_path)
+
+    def memget(self, idx):
+        self.getcount+=1
+        if self.getcount > self.refresh_flag and self.masking:
+            self.memcached_dataset = self.refresh_mask_ids()
+            self.getcount = 0
+        return self.memcached_dataset[idx]
+
+    def refresh_mask_ids(self):
+        self.logger.debug("Refreshing mask ids")
+        if self.masking:
+            self.mask_ids = self.build_mask_ids(self.input_length_cache)
+            
+            all_attention_mask = self.all_attention_mask.clone()
+            all_masklm = self.all_masklm.clone()
+
+            for idx in range(self.refresh_flag):
+                all_attention_mask[idx][self.mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+                all_masklm[idx][self.mask_ids[idx]] = self.all_input_ids[idx][self.mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+
+            return TensorDataset(self.all_input_ids, all_attention_mask, self.all_token_type_ids, all_masklm, self.all_lens, self.all_labels)
+        else:
+            return TensorDataset(self.all_input_ids, self.all_attention_mask, self.all_token_type_ids, self.all_masklm, self.all_lens, self.all_labels)
+
+    def build_mask_ids(self, input_length_cache):
+        # for each element, we get a set of indices that are randomly selected...
+        # Also, -2 and +1 take care of [cls] and [sep] not being masked
+        return [(torch.randperm(inplength-2)+1)[:int(inplength*self.mlm)]  for inplength in input_length_cache]
+
+    
+    def sharded_refresh_mask_ids(self):
+        pass
+
+    def sharded_build_mask_ids(self, input_length_cache):
+        pass
+
+
+    def convert_to_features(self, dataset, tokenizer, maxlen):
+        features = []
+        self.input_length_cache = []
+        for idx, sample in enumerate(dataset):
+            tokens = self.tokenizer.tokenize(sample[0])
+            if len(tokens) > maxlen - 2:
+                tokens = tokens[0:(maxlen - 2)]
+            # An easy way potentially -- get the index of each word in the token set...
+            # Then check which word is in keyword
+            # Then mask out that word...????????
+            finaltokens = ["[CLS]"]
+            token_type_ids = [0]
+            for token in tokens:
+                finaltokens.append(token)
+                token_type_ids.append(0)
+            finaltokens.append("[SEP]")
+            token_type_ids.append(0)
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(finaltokens)
+            attention_mask = [1]*len(input_ids)
+            input_len = len(input_ids)
+            self.input_length_cache.append(len(input_ids))
+            while len(input_ids) < maxlen:
+                input_ids.append(0)
+                attention_mask.append(0)
+                token_type_ids.append(0) 
+
+            assert len(input_ids) == maxlen
+            assert len(attention_mask) == maxlen
+            assert len(token_type_ids) == maxlen
+        
+            features.append(
+                (input_ids, attention_mask, token_type_ids, input_len, *sample[1:])
+            )
+
+        self.all_input_ids = torch.tensor([f[0] for f in features], dtype=torch.long)
+        self.all_attention_mask = torch.tensor([f[1] for f in features], dtype=torch.long)
+        self.all_token_type_ids = torch.tensor([f[2] for f in features], dtype=torch.long)
+        self.all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+        self.all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+        self.all_masklm = -1*torch.ones(self.all_input_ids.shape, dtype=torch.long)
+                
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.getter(idx)
+
+    def uncachedget(self, idx):
+        return 
+
+
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.generators import TextGenerator
+
+@edna.register_generator
+class FNCFilterMaskGenerator(TextGenerator):
+  # input includes tokenizer for build...
+
+  # Set it up such that, given a crawler, create a dataset from it.
+  # Then create a cached batch
+  # From cached batch, yield batches until it is empty
+
+    def build_transforms(self, transform, mode, **kwargs):  #<-- generator kwargs:
+        from ednaml.utils import locate_class
+        self.logger.info("Building Transforms")
+        tokenizer = kwargs.get("tokenizer", "HFAutoTokenizer")
+        self.tokenizer = locate_class(package="ednaml", subpackage="utils", classpackage=tokenizer, classfile="tokenizers")
+        self.tokenizer = self.tokenizer(**kwargs) 
+
+  
+    def buildDataset(self, crawler, mode, transform, **kwargs): #<-- dataset args:
+        self.logger.info("Building Dataset")
+        return FNCFilterMaskDataset(self.logger, crawler.metadata[mode]["crawl"], mode, tokenizer = self.tokenizer, crawler_secondary = crawler.metadata.get("secondary", {}), **kwargs) # needs maxlen, memcache, mlm_probability
+
+    def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
+        self.logger.info("Building Dataloader")
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size*self.gpus,
+                                        shuffle=kwargs.get("shuffle", True), num_workers = self.workers, 
+                                       collate_fn=self.collate_fn)
+
+    def getNumEntities(self, crawler, mode, **kwargs):  #<-- dataset args
+        label_dict = {
+            item: {"classes": crawler.metadata[mode]["classes"][item]}
+            for item in kwargs.get("classificationclass", ["color"])
+        }
+        return LabelMetadata(label_dict=label_dict)
+
+    def collate_fn(self, batch):
+        all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_labels  = map(torch.stack, zip(*batch))
+        max_len = max(all_lens).item()
+        all_input_ids = all_input_ids[:, :max_len]
+        all_attention_mask = all_attention_mask[:, :max_len]
+        all_token_type_ids = all_token_type_ids[:, :max_len]
+        all_masklm = all_masklm[:, :max_len]
+
+        return all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels
\ No newline at end of file
diff --git a/profiles/FNC/fnc-neighbors.py b/profiles/FNC/fnc-neighbors.py
new file mode 100644
index 00000000..c44e96fd
--- /dev/null
+++ b/profiles/FNC/fnc-neighbors.py
@@ -0,0 +1,44 @@
+import ednaml.core.decorators as edna
+
+
+
+from ednaml.deploy import BaseDeploy
+
+@edna.register_deployment
+class FNCPluginDeployment(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    prediction_scores, pooled_out, outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    
+    return prediction_scores, pooled_out, outputs
+  def output_setup(self, **kwargs):
+    neighbor_file = kwargs.get("neighbor_output")
+    unfilter_file = kwargs.get("unfiltered_output")
+    basename = kwargs.get("basename")
+    self.neighbor_name = "-".join([basename, neighbor_file])+".txt"
+    self.final_neighbor_name = "-".join([basename, neighbor_file])+".json"
+    self.unfilter_name = "-".join([basename, unfilter_file])+".json"
+
+    self.neighbor_obj = open(self.neighbor_name, "w")
+
+  def output_step(self, logits, features, secondary):
+    # outputs is a list, where the post is in outputs[2], I think????
+    # In any case, we need to check ouputs[2] to get the threshold...
+    write_list = (secondary[2]["FastKMP-l2"]["distance"] <= secondary[2]["FastKMP-l2"]["threshold"]).long().tolist()
+    self.neighbor_obj.write("\n".join([str(item) for item in write_list]) + "\n")
+
+  def end_of_deployment(self):
+    self.neighbor_obj.close()
+
+    nobj = open(self.neighbor_name, "r")
+    outputobj = open(self.final_neighbor_name, "w")
+    uobj = open(self.unfilter_name, "r")
+
+    fls = [nobj, uobj]
+    from itertools import zip_longest
+    for lines in zip_longest(*fls, fillvalue=""):
+      if int(lines[0].strip()) == 1:
+        outputobj.write(lines[1])
+
+    outputobj.close()
\ No newline at end of file
diff --git a/profiles/FNC/fnc-oracle.py b/profiles/FNC/fnc-oracle.py
new file mode 100644
index 00000000..67fca3eb
--- /dev/null
+++ b/profiles/FNC/fnc-oracle.py
@@ -0,0 +1,76 @@
+import ednaml.core.decorators as edna
+
+
+
+from ednaml.deploy import BaseDeploy
+
+@edna.register_deployment
+class FNCOracleBinning(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    prediction_scores, pooled_out, outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    
+    return prediction_scores, pooled_out, outputs
+  def output_setup(self, **kwargs):
+    oracle_file = kwargs.get("oracle_output")
+    extended_file = kwargs.get("extended_output")
+    basename = kwargs.get("basename")
+    self.oracle_name = "-".join([basename, oracle_file])+".txt"
+    self.final_oracle_name = "-".join([basename, oracle_file])+".json"
+    self.extended_name = "-".join([basename, extended_file])+".json"
+    self.target = kwargs.get("oracle_target", 500)
+    self.oracle_obj = open(self.oracle_name, "w")
+
+  def output_step(self, logits, features, secondary):
+    # outputs is a list, where the post is in outputs[2], I think????
+    # In any case, we need to check ouputs[2] to get the threshold...
+    dist = secondary[2]["FastKMP-l2"]["distance"].tolist()
+    idx = secondary[2]["FastKMP-l2"]["label"].tolist()
+    self.oracle_obj.write("\n".join([",".join([str(item[0]), str(item[1])]) for item in zip(dist, idx)]) + "\n")
+
+  def end_of_deployment(self):
+    self.oracle_obj.close()
+    
+
+    import pandas as pd
+    import math
+    import numpy as np
+    df = pd.read_csv(self.oracle_name, header=None, names=["distance", "label"])
+    df=df.reset_index()
+
+    oracle_per_proxy = math.ceil(self.target / len(df["label"].unique()))
+    proxies = df["label"].unique().tolist()
+
+    #for proxy in proxies:
+    extras = 0
+    indices = []
+    for proxy in proxies:
+      proxydf = df[df["label"] == proxy]
+      proxycount = proxydf.count()[0]
+      if proxycount < oracle_per_proxy:
+        proxy_select = proxycount
+        extras += oracle_per_proxy - proxycount
+      else:
+        proxy_select = oracle_per_proxy
+      proxy_quantiles = np.linspace(0,1,proxy_select, endpoint=False)
+      indices += proxydf.quantile(proxy_quantiles, interpolation="nearest")["index"].unique().tolist()
+      if len(indices) < proxy_select:
+        extras += proxy_select - len(indices)
+    sorted_indices = sorted(indices)
+
+    oracle_obj = open(self.final_oracle_name, "w")
+    extended_obj = open(self.extended_name, "r")
+    current_idx = 0
+    current_index = sorted_indices[current_idx]
+    for idx, line in enumerate(extended_obj):
+      if idx == current_index:
+        oracle_obj.write(line)
+        current_idx += 1
+        if current_idx >= len(sorted_indices):
+          break
+        current_index = sorted_indices[current_idx]
+      if idx%10000 == 0:
+        self.logger.debug("Processed %i points"%idx)
+    oracle_obj.close()
+    extended_obj.close()
\ No newline at end of file
diff --git a/profiles/FNC/labeling/config_base.yml b/profiles/FNC/labeling/config_base.yml
new file mode 100644
index 00000000..549edf84
--- /dev/null
+++ b/profiles/FNC/labeling/config_base.yml
@@ -0,0 +1,51 @@
+DEPLOYMENT:
+  DEPLOY: HFDeploy
+  FP16: False
+  EPOCHS: 1
+  OUTPUT_ARGS:
+    label_file_name: None
+  PLUGIN:
+    HOOKS: activated
+    RESET: False
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      azstorage: ednadatasets
+      azcontainer: edna-covid-extended # Adjust this to specific EFND dataset
+      azfile: tweets-2020-01-22.json.gz
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-extended-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: False
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [3] # This is the actual label...
+      annotation_idxs: [1,2]  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: unknown
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-expert-efnd
+  MODEL_BACKBONE: unknown
+  MODEL_QUALIFIER: cmu_miscov19
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 3 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
\ No newline at end of file
diff --git a/profiles/Fakeddit/fakeddit-v3.yml b/profiles/Fakeddit/fakeddit-v3.yml
index 45f360bd..64526592 100644
--- a/profiles/Fakeddit/fakeddit-v3.yml
+++ b/profiles/Fakeddit/fakeddit-v3.yml
@@ -1,7 +1,7 @@
 EXECUTION:
   DATAREADER:
     CRAWLER_ARGS: 
-      data_folder: Data/Fakeddit
+      data_folder: Data
     DATAREADER: AlbertReader
     DATASET_ARGS:
       classificationclass: ['reliability']
@@ -24,7 +24,7 @@ EXECUTION:
   TEST_FREQUENCY: 1
   TRAINER: BaseTrainer
   TRAINER_ARGS: 
-    accumulation_steps: 4
+    accumulation_steps: 8
 
 SAVE:
   MODEL_VERSION: 3
@@ -36,7 +36,7 @@ SAVE:
   CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/Fakeddit/Models/"
 
 TRANSFORMATION:
-  BATCH_SIZE: 128
+  BATCH_SIZE: 64
 
 MODEL:
   BUILDER: ednaml_model_builder
@@ -59,6 +59,7 @@ MODEL:
     initializer_range: 0.02
     layer_norm_eps: 1.0e-12
     pooling: pooled
+    num_labels: 6
   MODEL_NORMALIZATION: bn
   PARAMETER_GROUPS: [opt-1]
 
diff --git a/profiles/Fakeddit/main.py b/profiles/Fakeddit/fakeddit.py
similarity index 99%
rename from profiles/Fakeddit/main.py
rename to profiles/Fakeddit/fakeddit.py
index d4041cf0..32eb8b2b 100644
--- a/profiles/Fakeddit/main.py
+++ b/profiles/Fakeddit/fakeddit.py
@@ -1,6 +1,9 @@
 import ednaml, torch, os, csv
 from ednaml.crawlers import Crawler
 
+import ednaml.core.decorators as edna
+
+@edna.register_crawler
 class FakedditCrawler(Crawler):
     def __init__(self, logger = None, data_folder="Data/Fakeddit"):
         """Crawls the Data folder with all datasets already extracted to their individual folders.
@@ -59,6 +62,7 @@ def __init__(self, logger = None, data_folder="Data/Fakeddit"):
 from torch import nn
 from ednaml.utils.albert import AlbertEmbeddingAverage, AlbertPooledOutput, AlbertRawCLSOutput
 
+@edna.register_model
 class FakedditModel(ModelAbstract):
     def __init__(self, base, weights, metadata, normalization, parameter_groups, **kwargs):
         super().__init__(base=base,
@@ -155,6 +159,7 @@ def partial_load(self, weights_path):
 import tqdm
 import torch
 
+@edna.register_trainer
 class FakedditTrainer(BaseTrainer):
     def __init__(
             self,
diff --git a/profiles/MiDAS/midas-dataset.py b/profiles/MiDAS/midas-dataset.py
new file mode 100644
index 00000000..da67eae0
--- /dev/null
+++ b/profiles/MiDAS/midas-dataset.py
@@ -0,0 +1,222 @@
+import ednaml
+import csv, glob, os
+from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
+
+@edna.register_crawler
+class MiDASCrawler(Crawler):
+  def __init__(self, logger = None, data_folder="Data", include=["cmu_miscov19", "covid19_fn_title", "kagglefn_short"]):
+    """Crawls the Data folder with all datasets already extracted to their individual folders.
+    Assumes specific file construction: files are separated into splits 
+    (train, test, and val), with label subsets fake and true, with naming convention:
+
+    <datasetname>-<labelsubset>-<split>.csv
+
+    """
+    logger.info("Crawling %s for %s"%(data_folder, str(include)))
+    
+    self.classes = {}
+    self.classes["fnews"] = 2
+    self.metadata = {}
+
+    self.data_folder = data_folder
+    datasets_folders = glob.glob(os.path.join(data_folder, "*"))
+    datasets_folders = [item for item in datasets_folders if os.path.basename(item) in include]
+    self.classes["dataset"] = len(datasets_folders)
+
+
+    self.metadata["train"] = {}
+    self.metadata["test"] = {}
+    self.metadata["val"] = {}
+    self.metadata["train"]["crawl"] = []
+    self.metadata["test"]["crawl"] = []
+    self.metadata["val"]["crawl"] = []
+    for idx, folder in enumerate(datasets_folders):
+      ftrain = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "train.csv"])) for subset in ["fake", "true"]]
+      ftest = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "test.csv"])) for subset in ["fake", "true"]]
+      fval = [os.path.join(folder, "-".join([os.path.basename(folder), subset, "val.csv"])) for subset in ["fake", "true"]]
+
+      # train:
+      self.metadata["train"]["crawl"] += self.getTextAndLabels(ftrain, idx)
+      self.metadata["test"]["crawl"] += self.getTextAndLabels(ftest, idx)
+      self.metadata["val"]["crawl"] += self.getTextAndLabels(fval, idx)
+
+    self.metadata["train"]["classes"] = self.classes
+    self.metadata["test"]["classes"] = self.classes
+    self.metadata["val"]["classes"] = self.classes
+
+  def getTextAndLabels(self, listOfFiles, datasetidx):
+    crawl = []
+    for file in listOfFiles:
+      # For each file, open as csv. Extract the relevant columns ("text", "label")
+      with open(file, "r") as ofile:
+        csvread = csv.reader(ofile)
+
+        header = next(csvread)
+        textidx = [idx for idx,item in enumerate(header) if item == "text"][0]
+        labelidx = [idx for idx,item in enumerate(header) if item == "label"][0]
+
+        for row in csvread:
+          crawl.append([row[textidx], int(row[labelidx]), datasetidx])
+    return crawl
+
+
+
+import torch
+from torch.utils.data import TensorDataset
+class MiDASDataset(torch.utils.data.Dataset):
+  def __init__(self, dataset, mode, transform=None, **kwargs):
+    self.dataset = dataset  # list of tuples (text, label, datasetlabel)
+    self.cache = kwargs.get("cache", False)
+    self.memcache = kwargs.get("memcache", False)
+    self.tokenizer = kwargs.get("tokenizer")
+    self.maxlen = kwargs.get("maxlen")
+    self.mlm = kwargs.get("mlm_probability", 0.2)
+    self.masking = kwargs.get("masking", mode=="train")
+
+    self.getcount = 0
+    self.refresh_flag = len(dataset)
+
+    if self.cache and self.memcache:
+      raise ValueError("Use only one cache type")
+
+    self.getter = self.uncachedget
+    if self.cache:
+      raise NotImplementedError()
+      self.getter = self.diskget
+    if self.memcache:
+      self.getter = self.memget
+
+    # if cache, then we will cache transformations to disk, then load them
+    # if memcache, then we will cache to memory and load them
+
+    if self.cache or self.memcache:
+      # The actual cache-ing
+      self.input_length_cache = []
+      self.convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen)
+      self.memcached_dataset = self.refresh_mask_ids()
+      
+
+  def __len__(self):
+    return len(self.dataset)
+
+  def __getitem__(self, idx):
+    return self.getter(idx)
+
+  def uncachedget(self, idx):
+    return 
+
+  def memget(self, idx):
+    self.getcount+=1
+    if self.getcount > self.refresh_flag and self.masking:
+      self.memcached_dataset = self.refresh_mask_ids()
+      self.getcount = 0
+    return self.memcached_dataset[idx]
+
+  def refresh_mask_ids(self):
+    print("Refreshing mask ids")
+    if self.masking:
+      self.mask_ids = self.build_mask_ids(self.input_length_cache)
+      
+      all_attention_mask = self.all_attention_mask.clone()
+      all_masklm = self.all_masklm.clone()
+
+      for idx in range(self.refresh_flag):
+        all_attention_mask[idx][self.mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+        all_masklm[idx][self.mask_ids[idx]] = self.all_input_ids[idx][self.mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+
+      return TensorDataset(self.all_input_ids, all_attention_mask, self.all_token_type_ids, all_masklm, self.all_lens, self.all_labels, self.all_datalabels)
+    else:
+      return TensorDataset(self.all_input_ids, self.all_attention_mask, self.all_token_type_ids, self.all_masklm, self.all_lens, self.all_labels, self.all_datalabels)
+
+  def build_mask_ids(self, input_length_cache):
+    # for each element, we get a set of indices that are randomly selected...
+    # Also, -2 and +1 take care of [cls] and [sep] not being masked
+    return [(torch.randperm(inplength-2)+1)[:int(inplength*self.mlm)]  for inplength in input_length_cache]
+
+
+  def convert_to_features(self, dataset, tokenizer, maxlen):
+    features = []
+    self.input_length_cache = []
+    for idx, sample in enumerate(dataset):
+      tokens = self.tokenizer.tokenize(sample[0])
+      if len(tokens) > maxlen - 2:
+        tokens = tokens[0:(maxlen - 2)]
+
+      finaltokens = ["[CLS]"]
+      token_type_ids = [0]
+      for token in tokens:
+        finaltokens.append(token)
+        token_type_ids.append(0)
+      finaltokens.append("[SEP]")
+      token_type_ids.append(0)
+
+      input_ids = self.tokenizer.convert_tokens_to_ids(finaltokens)
+      attention_mask = [1]*len(input_ids)
+      input_len = len(input_ids)
+      self.input_length_cache.append(len(input_ids))
+      while len(input_ids) < maxlen:
+        input_ids.append(0)
+        attention_mask.append(0)
+        token_type_ids.append(0) 
+
+      assert len(input_ids) == maxlen
+      assert len(attention_mask) == maxlen
+      assert len(token_type_ids) == maxlen
+      
+      features.append(
+          (input_ids, attention_mask, token_type_ids, input_len, sample[1], sample[2])
+      )
+
+    self.all_input_ids = torch.tensor([f[0] for f in features], dtype=torch.long)
+    self.all_attention_mask = torch.tensor([f[1] for f in features], dtype=torch.long)
+    self.all_token_type_ids = torch.tensor([f[2] for f in features], dtype=torch.long)
+    self.all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+    self.all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+    self.all_datalabels = torch.tensor([f[5] for f in features], dtype=torch.long)
+    self.all_masklm = -1*torch.ones(self.all_input_ids.shape, dtype=torch.long)
+
+
+
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.generators import TextGenerator
+
+@edna.register_generator
+class MiDASGenerator(TextGenerator):
+  # input includes tokenizer for build...
+
+  # Set it up such that, given a crawler, create a dataset from it.
+  # Then create a cached batch
+  # From cached batch, yield batches until it is empty
+
+  def build_transforms(self, transform, mode, **kwargs):  #<-- generator kwargs:
+    from ednaml.utils import locate_class
+    tokenizer = kwargs.get("tokenizer", "AlbertFullTokenizer")
+    self.tokenizer = locate_class(package="ednaml", subpackage="utils", classpackage=tokenizer, classfile="tokenizers")
+    self.tokenizer = self.tokenizer(**kwargs) # vocab_file, do_lower_case, spm_model_file
+
+  
+  def buildDataset(self, datacrawler, mode, transform, **kwargs): #<-- dataset args:
+    return MiDASDataset(datacrawler.metadata[mode]["crawl"], mode, tokenizer = self.tokenizer, **kwargs) # needs maxlen, memcache, mlm_probability
+
+  def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
+    return torch.utils.data.DataLoader(dataset, batch_size=batch_size*self.gpus,
+                                        shuffle=True, num_workers = self.workers, 
+                                       collate_fn=self.collate_fn)
+
+  def getNumEntities(self, crawler, mode, **kwargs):  #<-- dataset args
+    label_dict = {
+        item: {"classes": crawler.metadata[mode]["classes"][item]}
+        for item in kwargs.get("classificationclass", ["color"])
+    }
+    return LabelMetadata(label_dict=label_dict)
+
+  def collate_fn(self, batch):
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_labels, all_datalabels  = map(torch.stack, zip(*batch))
+    max_len = max(all_lens).item()
+    all_input_ids = all_input_ids[:, :max_len]
+    all_attention_mask = all_attention_mask[:, :max_len]
+    all_token_type_ids = all_token_type_ids[:, :max_len]
+    all_masklm = all_masklm[:, :max_len]
+
+    return all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels, all_datalabels
\ No newline at end of file
diff --git a/profiles/MiDAS/midas-encoder.yml b/profiles/MiDAS/midas-encoder.yml
index c3917a9b..8e31f294 100644
--- a/profiles/MiDAS/midas-encoder.yml
+++ b/profiles/MiDAS/midas-encoder.yml
@@ -3,7 +3,7 @@ EXECUTION:
     CRAWLER_ARGS: 
       data_folder: Data
       include: ["cmu_miscov19", "cov19_fn_title", "kagglefn_short", "cov_rumor"]
-    DATAREADER: DataReader
+    DATAREADER: AlbertReader
     DATASET_ARGS:
       classificationclass: ['fnews', 'dataset']
       maxlen: 512
diff --git a/profiles/MiDAS/midas-expert.py b/profiles/MiDAS/midas-expert.py
new file mode 100644
index 00000000..f3143fb9
--- /dev/null
+++ b/profiles/MiDAS/midas-expert.py
@@ -0,0 +1,96 @@
+import torch
+from ednaml.models.Albert import AlbertPreTrainedModel, AlbertModel, AlbertOnlyMLMHead
+from ednaml.models.Albert import AlbertConfig
+from torch.nn import CrossEntropyLoss
+from ednaml.models.ModelAbstract import ModelAbstract
+from torch import nn
+
+from ednaml.utils.albert import AlbertEmbeddingAverage, AlbertPooledOutput, AlbertRawCLSOutput
+import ednaml.core.decorators  as edna
+
+@edna.register_model
+class MiDASExpert(ModelAbstract):
+  def __init__(self, base, weights, metadata, normalization, parameter_groups, **kwargs):
+    super().__init__(base=base,
+          weights=weights,
+          metadata=metadata, 
+          normalization=normalization,
+          parameter_groups=parameter_groups,
+          **kwargs)
+    
+  def model_attributes_setup(self, **kwargs):
+    self.config = AlbertConfig(**kwargs)
+    self.configargs = self.config.getVars()
+    self.num_labels = kwargs.get("num_classes")
+    self.pool_method = kwargs.get("pooling", "pooled")
+  
+  def model_setup(self, **kwargs):
+    
+    self.dropout = nn.Dropout(0.1 if self.config.hidden_dropout_prob == 0 else self.config.hidden_dropout_prob)
+    self.classifier = nn.Linear(self.config.hidden_size, self.config.num_labels)
+
+    if self.pool_method == "pooled":
+      self.pooler_layer = AlbertPooledOutput()
+    elif self.pool_method == "raw":
+      self.pooler_layer = AlbertRawCLSOutput()
+    elif self.pool_method == "average":
+      self.pooler_layer = AlbertEmbeddingAverage()
+    else:
+      raise NotImplementedError()
+
+    self.init_weights()
+    
+    self.encoder, errors = AlbertModel.from_pretrained("pytorch_model.bin", config=self.config, output_loading_info=True)
+    print("Errors \n\t", errors)
+    
+        
+  def _tie_or_clone_weights(self, first_module, second_module):
+        """ Tie or clone module weights depending of weither we are using TorchScript or not
+        """
+
+        if self.config.torchscript:
+            first_module.weight = nn.Parameter(second_module.weight.clone())
+        else:
+            first_module.weight = second_module.weight
+
+
+        if hasattr(first_module, 'bias') and first_module.bias is not None:
+            first_module.bias.data = torch.nn.functional.pad(
+                first_module.bias.data,
+                (0, first_module.weight.shape[0] - first_module.bias.shape[0]),
+                'constant',
+                0
+            )
+
+  def forward_impl(self, x, attention_mask=None, token_type_ids=None, position_ids=None, head_mask=None):
+    outputs = self.encoder(x,
+                        attention_mask=attention_mask,
+                        token_type_ids=token_type_ids,
+                        position_ids=position_ids,
+                        head_mask=head_mask)  # sequence_output, pooled_output, (hidden_states), (attentions)
+    
+    pooled_output = self.pooler_layer(outputs)  # TODO -- have an option to either use the pooled output, the original output, or average the embeddings together, i.e. a layer that is either a lambda layer, or does some averaging...
+    pooled_output = self.dropout(pooled_output+0.1)
+    logits = self.classifier(pooled_output)
+    return logits, pooled_output, outputs[2:] # list of k scores; hidden states, attentions...
+
+  def init_weights(self):
+        """ Initialize and prunes weights if needed. """
+        # Initialize weights
+        self.apply(self._init_weights)
+
+  def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, torch.nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+  def partial_load(self, weights_path):
+    super().partial_load(self, weights_path)  # For this, we need to look at the from_pretrained function to accurately load the saved weights from .bin...
\ No newline at end of file
diff --git a/profiles/MiDAS/midas-expert.yml b/profiles/MiDAS/midas-expert.yml
index 97273a77..2c1e46aa 100644
--- a/profiles/MiDAS/midas-expert.yml
+++ b/profiles/MiDAS/midas-expert.yml
@@ -3,7 +3,7 @@ EXECUTION:
     CRAWLER_ARGS: 
       data_folder: Data
       include: ["cmu_miscov19"]
-    DATAREADER: DataReader
+    DATAREADER: AlbertReader
     DATASET_ARGS:
       classificationclass: ['fnews', 'dataset']
       maxlen: 512
diff --git a/profiles/MiDAS/midas-trainer.py b/profiles/MiDAS/midas-trainer.py
new file mode 100644
index 00000000..53865963
--- /dev/null
+++ b/profiles/MiDAS/midas-trainer.py
@@ -0,0 +1,132 @@
+import ednaml
+from ednaml.models.ModelAbstract import ModelAbstract
+import torch
+from ednaml.crawlers import Crawler
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.trainer import BaseTrainer
+from torch.utils.data import DataLoader
+from logging import Logger
+from typing import Dict, List
+from ednaml.config.EdnaMLConfig import EdnaMLConfig
+import numpy as np
+from sklearn.metrics import f1_score
+import tqdm
+
+import ednaml.core.decorators as edna
+
+@edna.register_trainer
+class MiDASExpertTrainer(BaseTrainer):
+  def __init__(
+        self,
+        model: ModelAbstract,
+        loss_fn: List[ednaml.loss.builders.LossBuilder],
+        optimizer: List[torch.optim.Optimizer],
+        loss_optimizer: List[torch.optim.Optimizer],
+        scheduler: List[torch.optim.lr_scheduler._LRScheduler],
+        loss_scheduler: List[torch.optim.lr_scheduler._LRScheduler],
+        train_loader: DataLoader,
+        test_loader: DataLoader,
+        epochs: int,
+        skipeval: bool,
+        logger: Logger,
+        crawler: Crawler,
+        config: EdnaMLConfig,
+        labels: LabelMetadata,
+        **kwargs
+    ):
+    super().__init__(
+            model,
+            loss_fn,
+            optimizer,
+            loss_optimizer,
+            scheduler,
+            loss_scheduler,
+            train_loader,
+            test_loader,
+            epochs,
+            skipeval,
+            logger,
+            crawler,
+            config,
+            labels,
+            **kwargs
+        )
+    self.softaccuracy = []
+
+  def step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels, all_datalabels = batch
+    outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+
+    logits = outputs[0]
+    logits_loss = self.loss_fn["classification"](
+        logits=logits, labels=all_labels
+    )
+    softmax_accuracy = (
+                (logits.max(1)[1] == all_labels)
+                .float()
+                .mean()
+            )
+
+    self.losses["classification"].append(logits_loss.item())
+    self.softaccuracy.append(softmax_accuracy.cpu().item())
+    
+    return logits_loss
+
+  def evaluate_impl(self):
+    logits, labels, dlabels = [],[],[]
+    with torch.no_grad():
+      for batch in tqdm.tqdm(
+          self.test_loader, total=len(self.test_loader), leave=False
+      ):
+        batch = tuple(item.cuda() for item in batch)
+        all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels, all_datalabels = batch
+        outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+        logit = outputs[0].detach().cpu()
+        label = all_labels.detach().cpu()
+        dlabel = all_datalabels.detach().cpu()
+
+        logits.append(logit)
+        labels.append(label)
+        dlabels.append(dlabel)
+
+
+        #accuracy.append(
+        #    (torch.argmax(logits.detach().cpu(), dim=1) == all_labels.detach().cpu()).sum() / all_labels.size(0)
+        #)
+    logits, labels, dlabels = (
+        torch.cat(logits, dim=0),
+        torch.cat(labels, dim=0),
+        torch.cat(dlabels, dim=0)
+    )
+    self.logger.info("Obtained logits and labels, validation in progress")
+
+
+    logit_labels = torch.argmax(logits, dim=1)
+    accuracy = (logit_labels == labels).sum().float() / float(labels.size(0))
+    micro_fscore = np.mean(f1_score(labels, logit_labels, average="micro"))
+    weighted_fscore = np.mean(f1_score(labels, logit_labels, average="weighted"))
+    self.logger.info("\tAccuracy: {:.3%}".format(accuracy))
+    self.logger.info("\tMicro F-score: {:.3f}".format(micro_fscore))
+    self.logger.info("\tWeighted F-score: {:.3f}".format(weighted_fscore))
+
+    return logit_labels, (labels, dlabels), logits
+
+
+  def printStepInformation(self):
+        loss_avg = 0.0
+        for lossname in self.losses:
+            loss_avg += (
+                sum(self.losses[lossname][-self.step_verbose :])
+                / self.step_verbose
+            )
+        #loss_avg /= self.num_losses
+        soft_avg = sum(self.softaccuracy[-100:]) / float(
+            len(self.softaccuracy[-100:])
+        )
+
+        self.logger.info(
+            "Epoch{0}.{1}\tClassification Loss: {2:.3f}\t Training Accuracy:  {3:.3f}".format(
+                self.global_epoch, self.global_batch, loss_avg, soft_avg
+                )
+        )
diff --git a/profiles/NELA/deploy-plugin-generation.py b/profiles/NELA/deploy-plugin-generation.py
new file mode 100644
index 00000000..732df77f
--- /dev/null
+++ b/profiles/NELA/deploy-plugin-generation.py
@@ -0,0 +1,25 @@
+import ednaml, torch, os, csv
+from ednaml.deploy.BaseDeploy import BaseDeploy
+import ednaml.core.decorators as edna
+
+@edna.register_deployment
+class NELADeploy(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    logits, features, secondary = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    logit = logits.detach().cpu()
+    logit_labels = torch.argmax(logit, dim=1)
+    #import pdb
+    #pdb.set_trace()
+    #print(str(logit_labels[0]), "\t", ",".join([str(item) for item in all_input_ids[0].cpu()[:10].tolist()]))
+    if len(secondary)>1:
+      plugin_post = secondary[2]
+    else:
+      plugin_post = None
+    return logit_labels, logit, plugin_post
+
+  def output_setup(self, **kwargs):
+    pass
+  def output_step(self, logits, features, secondary):
+    pass
\ No newline at end of file
diff --git a/profiles/NELA/deploy-plugin-usage.py b/profiles/NELA/deploy-plugin-usage.py
new file mode 100644
index 00000000..431af0da
--- /dev/null
+++ b/profiles/NELA/deploy-plugin-usage.py
@@ -0,0 +1,79 @@
+import ednaml, torch, os, csv
+from ednaml.crawlers import Crawler
+from ednaml.deploy.BaseDeploy import BaseDeploy
+import ednaml.core.decorators as edna
+
+
+# For each batch, get the predicted labels, and save the labels in an append text file
+
+@edna.register_deployment
+class NELADeploy(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    logits, features, secondary = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    logit = logits.detach().cpu()
+    logit_labels = torch.argmax(logit, dim=1)
+    #import pdb
+    #pdb.set_trace()
+    #print(str(logit_labels[0]), "\t", ",".join([str(item) for item in all_input_ids[0].cpu()[:10].tolist()]))
+    if len(secondary)>1:
+      plugin_post = secondary[2]
+    else:
+      plugin_post = None
+    return logit_labels, logit, plugin_post
+
+  def output_setup(self, **kwargs):
+    file_path = kwargs.get("file_name", "deploy-out.txt")
+    self.fobj = open(file_path, "w")
+
+  def output_step(self, logits, features, secondary):
+    predicted_label = logits.tolist()
+    l2_predicted_label = secondary["KMP-l2"]["cluster_mean"].tolist()
+    l2_dist = secondary["KMP-l2"]["distance"].to(torch.float32).tolist()
+    cos_predicted_label = secondary["KMP-cos"]["cluster_mean"].tolist()
+    cos_dist = secondary["KMP-cos"]["distance"].to(torch.float32).tolist()
+    l_score = secondary["RL-midas"]["l_score"]
+    smooth_l_score = secondary["RL-midas"]["smooth_l_score"]
+    l_threshold = secondary["RL-midas"]["l_threshold"]
+    smooth_l_threshold = secondary["RL-midas"]["smooth_l_threshold"]
+
+    output_list = [",".join(map(str,item)) for item in zip(
+      predicted_label,
+      l2_predicted_label,
+      l2_dist,
+      cos_predicted_label,
+      cos_dist,
+      l_score,
+      smooth_l_score,
+      l_threshold,
+      smooth_l_threshold
+      )]
+    self.fobj.write("\n".join(output_list)+"\n")
+
+#@edna.register_deployment
+class NELADeployPluginSetup(NELADeploy):
+  def output_setup(self, **kwargs):
+    pass
+  def output_step(self, logits, features, secondary):
+    pass
+
+
+
+import click
+@click.argument("config")
+@click.argument("mode")
+def main(config, deploy, mode):
+    from ednaml.core import EdnaDeploy
+    ed = EdnaDeploy(config=config, deploy=deploy)
+    ed.add("./nela.py")
+    ed.addModelClass(main.NELAModel)
+    ed.addDeploymentClass(NELADeploy)
+
+    ed.apply(input_size=(ed.cfg.TRAIN_TRANSFORMATION.BATCH_SIZE,ed.cfg.EXECUTION.DATAREADER.DATASET_ARGS["maxlen"]),
+          dtypes=[torch.long])
+
+    ed.deploy()
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/profiles/NELA/deploy.py b/profiles/NELA/deploy.py
deleted file mode 100644
index 0b799701..00000000
--- a/profiles/NELA/deploy.py
+++ /dev/null
@@ -1,41 +0,0 @@
-import ednaml, torch, os, csv
-from ednaml.crawlers import Crawler
-from ednaml.deploy.BaseDeploy import BaseDeploy
-
-
-
-
-
-class NELADeploy(BaseDeploy):
-
-
-    def deploy_step(self, batch):
-        return super().deploy_step(batch)
-
-    def output_setup(self, **kwargs):
-        return super().output_setup(**kwargs)
-
-    def output_step(self, logits, features, secondary):
-        return super().output_step(logits, features, secondary)
-
-
-
-
-import click
-@click.argument("config")
-@click.argument("mode")
-def main(config, mode):
-    from ednaml.core import EdnaDeploy
-    import main
-    eml = EdnaDeploy(config=config)
-    eml.addCrawlerClass(main.NELACrawler)
-    eml.addModelClass(main.NELAModel)
-    eml.addDeploymentClass(NELADeploy)
-
-    eml.apply(input_size=(eml.cfg.TRAIN_TRANSFORMATION.BATCH_SIZE,eml.cfg.EXECUTION.DATAREADER.DATASET_ARGS["maxlen"]),
-          dtypes=[torch.long])
-
-    eml.train()
-
-if __name__ == "__main__":
-    main()
\ No newline at end of file
diff --git a/profiles/NELA/nela-covid-deploy.yml b/profiles/NELA/nela-covid-deploy.yml
index d9437acc..8183bef8 100644
--- a/profiles/NELA/nela-covid-deploy.yml
+++ b/profiles/NELA/nela-covid-deploy.yml
@@ -16,6 +16,7 @@ DEPLOYMENT:
       mlm_probability: 0
       masking: False
       shuffle: False
+      data_shuffle: False
     GENERATOR_ARGS: 
       do_lower_case: True
       spm_model_file: 30k-clean.model 
diff --git a/profiles/NELA/main.py b/profiles/NELA/nela.py
similarity index 98%
rename from profiles/NELA/main.py
rename to profiles/NELA/nela.py
index 812ae936..6f9b45dd 100644
--- a/profiles/NELA/main.py
+++ b/profiles/NELA/nela.py
@@ -1,6 +1,9 @@
-import ednaml, torch, os, csv
+import ednaml, torch, os, csv, os, json, glob, click
 from ednaml.crawlers import Crawler
+import ednaml.core.decorators as edna
 
+
+@edna.register_crawler
 class NELACrawler(Crawler):
   def __init__(self, logger = None, data_folder="Data", sub_folder="nela-covid-2020"):
     """Crawls the Data folder with all datasets already extracted to their individual folders.
@@ -98,7 +101,7 @@ def __init__(self, logger = None, data_folder="Data", sub_folder="nela-covid-202
 from torch import nn
 
 from ednaml.utils.albert import AlbertEmbeddingAverage, AlbertPooledOutput, AlbertRawCLSOutput
-
+@edna.register_model
 class NELAModel(ModelAbstract):
   def __init__(self, base, weights, metadata, normalization, parameter_groups, **kwargs):
     super().__init__(base=base,
@@ -196,6 +199,7 @@ def partial_load(self, weights_path):
 import tqdm
 import torch
 
+@edna.register_trainer
 class NELATrainer(BaseTrainer):
   def __init__(
         self,
diff --git a/profiles/NELA/plugins.yml b/profiles/NELA/plugins.yml
new file mode 100644
index 00000000..445dc2a3
--- /dev/null
+++ b/profiles/NELA/plugins.yml
@@ -0,0 +1,22 @@
+MODEL_PLUGIN:
+  - PLUGIN: KMeansProxy
+    PLUGIN_NAME: KMP-l2
+    PLUGIN_KWARGS:
+      num_clusters: 20
+      dimensions: 768
+      dist: euclidean
+      epochs: 5
+  - PLUGIN: KMeansProxy
+    PLUGIN_NAME: KMP-cos
+    PLUGIN_KWARGS:
+      num_clusters: 20
+      dimensions: 768
+      dist: cosine
+      epochs: 5
+  - PLUGIN: RandomizedLipschitz
+    PLUGIN_NAME: RL-midas
+    PLUGIN_KWARGS:
+      proxies: 15
+      dimensions: 768
+      dist: euclidean
+      proxy_epochs: 5
diff --git a/profiles/PolitiFact/config.yml b/profiles/PolitiFact/config.yml
new file mode 100644
index 00000000..a7e43e60
--- /dev/null
+++ b/profiles/PolitiFact/config.yml
@@ -0,0 +1,117 @@
+EXECUTION:
+  EPOCHS: 10
+  FP16: False
+  MODEL_SERVING: False
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  SKIPEVAL: False
+  TEST_FREQUENCY: 1
+  TRAINER: HFTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 8
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      split: 0.9
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-politifact-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [1] # This is the actual label...
+      annotation_idxs: []  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+
+DEPLOYMENT:
+  DEPLOY: HFDeploy
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: warmup
+    RESET: False
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      split: 0.9
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-politifact-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [1] # This is the actual label...
+      annotation_idxs: []  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+    
+LOSS: 
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    LABEL: fnews
+    NAME: classification
+
+OPTIMIZER:
+  - OPTIMIZER: AdamW
+    BASE_LR: 1.0e-5
+    LR_BIAS_FACTOR: 1.0
+    OPTIMIZER_KWARGS: 
+      eps: 1.0e-6
+    OPTIMIZER_NAME: opt-1
+    WEIGHT_BIAS_FACTOR: 0.0005
+    WEIGHT_DECAY: 0.0005
+
+SCHEDULER:
+  - LR_KWARGS: 
+      step_size: 5
+      gamma: 0.5
+    LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-politifact
+  MODEL_BACKBONE: albert
+  MODEL_QUALIFIER: base
+  DRIVE_BACKUP: False
+  LOG_BACKUP: False
+  SAVE_FREQUENCY: 2 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: HFAutoModel
+  MODEL_BASE: albert-base-v2
+  MODEL_KWARGS: 
+    auto_class: AutoModelForSequenceClassification
+    hidden_act: gelu
+    pooling: pooled
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
\ No newline at end of file
diff --git a/profiles/PolitiFact/gossipcop.yml b/profiles/PolitiFact/gossipcop.yml
new file mode 100644
index 00000000..5bb8d037
--- /dev/null
+++ b/profiles/PolitiFact/gossipcop.yml
@@ -0,0 +1,125 @@
+EXECUTION:
+  EPOCHS: 10
+  FP16: False
+  MODEL_SERVING: False
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  SKIPEVAL: True
+  TEST_FREQUENCY: 1
+  TRAINER: HFTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 8
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      split: 0.9
+      true_url: https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/gossipcop_real.csv
+      fake_url: https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/gossipcop_fake.csv
+      true_file: gossipcop_real.csv
+      fake_file: gossipcop_fake.csv
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-gossipcop-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [1] # This is the actual label...
+      annotation_idxs: []  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+
+DEPLOYMENT:
+  DEPLOY: HFDeploy
+  FP16: False
+  EPOCHS: 1
+  PLUGIN:
+    HOOKS: warmup
+    RESET: False
+  DATAREADER:
+    DATAREADER: HFReader
+    CRAWLER_ARGS: 
+      split: 0.9
+      true_url: https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/gossipcop_real.csv
+      fake_url: https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/gossipcop_fake.csv
+      true_file: gossipcop_real.csv
+      fake_file: gossipcop_fake.csv
+    DATASET_ARGS:
+      classificationclass: ['fnews']
+      maxlen: 512
+      shardcache: True
+      shardsize: 20000
+      shard_replace: False
+      shardpath: datashard-artifacts
+      shardname: fnc-gossipcop-shard
+      mlm_probability: 0.15
+      shuffle: False
+      data_shuffle: True
+      masking: False
+      keyword_mask: True
+      keywords: []
+      label_idxs: [1] # This is the actual label...
+      annotation_idxs: []  # this is shares and likes
+    GENERATOR_ARGS: 
+      tokenizer: HFAutoTokenizer
+      from_pretrained: albert-base-v2
+    
+LOSS: 
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    LABEL: fnews
+    NAME: classification
+
+OPTIMIZER:
+  - OPTIMIZER: AdamW
+    BASE_LR: 1.0e-5
+    LR_BIAS_FACTOR: 1.0
+    OPTIMIZER_KWARGS: 
+      eps: 1.0e-6
+    OPTIMIZER_NAME: opt-1
+    WEIGHT_BIAS_FACTOR: 0.0005
+    WEIGHT_DECAY: 0.0005
+
+SCHEDULER:
+  - LR_KWARGS: 
+      step_size: 5
+      gamma: 0.5
+    LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+
+LOGGING:
+  STEP_VERBOSE: 100
+  INPUT_SIZE: [16, 512]
+
+TRANSFORMATION:
+  BATCH_SIZE: 16
+  WORKERS: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: fnc-gossipcop
+  MODEL_BACKBONE: albert
+  MODEL_QUALIFIER: base
+  DRIVE_BACKUP: False
+  LOG_BACKUP: False
+  SAVE_FREQUENCY: 2 # Epoch
+  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/FNC/Models/"
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: HFAutoModel
+  MODEL_BASE: albert-base-v2
+  MODEL_KWARGS: 
+    auto_class: AutoModelForSequenceClassification
+    hidden_act: gelu
+    pooling: pooled
+  MODEL_NORMALIZATION: bn
+  PARAMETER_GROUPS: [opt-1]
\ No newline at end of file
diff --git a/profiles/PolitiFact/politifact-crawler.py b/profiles/PolitiFact/politifact-crawler.py
new file mode 100644
index 00000000..097b3129
--- /dev/null
+++ b/profiles/PolitiFact/politifact-crawler.py
@@ -0,0 +1,76 @@
+import ednaml, torch, csv, ctypes as ct, os
+import ednaml.core.decorators as edna
+from ednaml.crawlers import Crawler
+from ednaml.utils.web import download
+
+@edna.register_crawler
+class PolitifactCrawler(Crawler):
+    def __init__(self, logger = None, split = 0.9, true_url = None, fake_url = None, true_file = "politifact_real.csv", fake_file = "politifact_fake.csv"):
+        csv.field_size_limit(int(ct.c_ulong(-1).value // 2))
+        self.true_url = true_url
+        self.fake_url = fake_url
+        
+        if self.fake_url is None:
+            self.fake_url = "https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/politifact_fake.csv"
+        if self.true_url is None:
+            self.true_url = "https://raw.githubusercontent.com/KaiDMML/FakeNewsNet/master/dataset/politifact_real.csv"
+        
+        self.fake_file = fake_file
+        self.true_file = true_file
+        
+        if not os.path.exists(self.fake_file):
+            download(self.fake_file, self.fake_url)
+        if not os.path.exists(self.true_file):
+            download(self.true_file, self.true_url)
+
+        all_fakes = []
+        all_true = []
+
+
+        with open(self.fake_file, "r") as ffile:
+            fobj = csv.reader(ffile, delimiter=",")
+            header = next(fobj)
+            for row in fobj:
+                all_fakes.append(
+                    (row[2], 0)
+                )
+        
+        with open(self.true_file, "r") as ffile:
+            fobj = csv.reader(ffile, delimiter=",")
+            header = next(fobj)
+            for row in fobj:
+                all_true.append(
+                    (row[2], 1)
+                )
+
+        import random
+        random.seed(345345)
+        random.shuffle(all_fakes)
+        random.shuffle(all_true)
+        
+        self.metadata = {}
+        self.metadata["train"] = {}
+        self.metadata["test"] = {}
+        self.metadata["val"] = {}
+        self.metadata["train"]["crawl"] = []
+        self.metadata["test"]["crawl"] = []
+        self.metadata["val"]["crawl"] = []
+        
+
+        fsplit = int(len(all_fakes)*split)
+        tsplit = int(len(all_true)*split)
+        self.metadata["train"]["crawl"] += all_fakes[:fsplit]
+        self.metadata["train"]["crawl"] += all_true[:tsplit]
+        self.metadata["test"]["crawl"] += all_fakes[fsplit:]
+        self.metadata["test"]["crawl"] += all_true[tsplit:]
+
+        random.shuffle(self.metadata["train"]["crawl"])
+        random.shuffle(self.metadata["test"]["crawl"])
+
+        self.classes = {}
+        self.classes["fnews"] = 2
+        
+        self.metadata["train"]["classes"] = self.classes
+        self.metadata["test"]["classes"] = self.classes
+
+        
\ No newline at end of file
diff --git a/readme.md b/readme.md
index f6c32b0e..3596b3ba 100644
--- a/readme.md
+++ b/readme.md
@@ -1,9 +1,4 @@
-# CoLabel
+# EdnaML
 
 
-This repo contains the code for 'Constructive Interpretability with CoLabel: Corroborative Integration, Complementary Features, and Collaborative Learning'
-
-
-# Requirements
-The code uses PyTorch, as well as existing vehicle datasets.
-
+This repo contains the EdnaML Framework.
diff --git a/roadmap.md b/roadmap.md
index 2835b56b..b45c3abe 100644
--- a/roadmap.md
+++ b/roadmap.md
@@ -1,3 +1,7 @@
+🟢🟡
+
+Deploy also needs to deal with storage classes
+
 Roadmap
 [
 
@@ -24,8 +28,8 @@ Else, ignore and write -1
 
 
 - Deploy
-    - need a way to rerun existing saved model, by also saving the file model was generated from...
-    - for actual file, pretty easy
+    - need a way to rerun existing saved model, by also saving the file model was generated from...🟢
+    - for actual file, pretty easy🟢
     - but for colab, maybe we can have an edna API that takes in the file...? Yeah -- that's it...
 
 
@@ -144,7 +148,7 @@ Else, ignore and write -1
 - VehicleColors 🟢
 - logging -- tensorboard???
 - metrics logging
-- apply()
+- apply()🟢
     - also, when user wants to put custom things in Edna, they can use the custom add functions (addOptimizer, etc). This will do lazy setup
     - Then, when user does apply, this will apply the lazy changes internally.
         - custom Crawler    🟢
@@ -169,7 +173,7 @@ Else, ignore and write -1
     - need to upgrade to azure, or in-house system, or cloudlab...
 
 
-- EdnaML in test mode -- does not build train data loader...
+- EdnaML in test mode -- does not build train data loader...🟢
 --------------------------------------------------------------------
 so, someone does the following??? ✅
 EdnaML() <-- empty, so no configuration is set up> ✅
@@ -194,11 +198,11 @@ add/update LossScheduler ✅
 
 
 --------------------------------------------------------------------
-soft_targets: true
-soft_target_branch: ['vcolor', 'vtype']
-soft_target_output_source: fuse
+soft_targets: true🟢
+soft_target_branch: ['vcolor', 'vtype']🟢
+soft_target_output_source: fuse🟢
 
-if soft targets true, model internally stores this value as well as the name of the output source. Then, secondary outputs gets a list of these soft target `guesses`
+if soft targets true, model internally stores this value as well as the name of the output source. Then, secondary outputs gets a list of these soft target `guesses`🟢
 --------------------------------------------------------------------
 
 
@@ -452,4 +456,14 @@ for batch in dataloader:
 
     get these values
     compare to true ground truth... 
-    
\ No newline at end of file
+
+
+
+
+deployment_step:🟢
+
+    for batch in data:🟢
+        <perform the live KNN proxy thingamajig>🟢
+        <pass it though the model to get the L score stuff>🟢
+        update the components with the right values, i.e. the KMeans matrix, and the L-score for each proxy as well...(i.e. we might need multiple epochs for this)🟢
+
diff --git a/setup.cfg b/setup.cfg
index e51577e2..27cc2db1 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -32,5 +32,13 @@ include_package_data = true
 python_requires = >= 3.7
 # Dependencies are in setup.py for GitHub's dependency graph.
 
-[options.packages.find]
-where = src
+[options.extras_require]
+azure = 
+    azure-storage-blob >= 12.13.0
+metrics = 
+    torchmetrics >= 0.9.3
+transformers = 
+    transformers >= 4.21.*
+full = ednaml[azure,metrics]
+
+
diff --git a/setup.py b/setup.py
index 1969decf..f1a4e691 100644
--- a/setup.py
+++ b/setup.py
@@ -9,7 +9,8 @@
         "torchvision>=0.11.*",
         "Pillow>=7.1.2",    #9.0.x, this is for colab...
         "tqdm>=4.63.*",
-        "sentencepiece>=0.1.96"
+        "sentencepiece>=0.1.96",
+        "sortedcontainers>=2.4.0"
     ]
 
 
@@ -19,4 +20,7 @@
     name="ednaml",
     install_requires=requires,
     keywords=["deep learning", "pytorch", "AI"],
-)
\ No newline at end of file
+)
+
+# Azure :
+# Install azure-storage>=0.30.0
\ No newline at end of file
diff --git a/src/ednaml/backbones/multibranchresnet.py b/src/ednaml/backbones/multibranchresnet.py
index 0adfcd50..0c85fca7 100644
--- a/src/ednaml/backbones/multibranchresnet.py
+++ b/src/ednaml/backbones/multibranchresnet.py
@@ -36,7 +36,7 @@ class multibranchresnet(nn.Module):
         part_attention (bool): Whether the model uses local/part attention at the first resnet block
         secondary_attention (None | int): Whether `attention` is applied to all blocks (None) or to the specified block (int).
         shared_block_count (int): Number of resnet blocks with weight sharing. Maximum value 4
-        num_branches (int): Number of branches after weight-shared blocks
+        number_branches (int): Number of branches after weight-shared blocks
         pytorch_weights_paths (Dict[str,int]): Strings corresponding to official imagenet resnet weights from pytorch
         resnetinput (ResnetInput): The input block consisting of conv layer, relu, and pooling.
         sharedblock (Union[nn.Sequential,nn.Identity]): The shared layers, consisting of at most 4 Resnet blocks.
@@ -75,7 +75,7 @@ def __init__(
         secondary_attention: int = None,
         ia_attention: bool = None,
         part_attention: bool = None,
-        num_branches: int = 2,
+        number_branches: int = 2,
         shared_block: int = 0,
         **kwargs
     ):
@@ -207,7 +207,7 @@ def __init__(
         # Then, given the branches, put the remaining resnet blocks in their branches
         # During prediction, we will just get branch features so order does not matter yet. it will matter in MultiBranchResnet
         # So, self.branches will be a nn.moduleList, with a bunch of nn.Sequentials
-        self.num_branches = num_branches
+        self.num_branches = number_branches
         branches = [None] * self.num_branches
         for bidx in range(self.num_branches):
             branches[bidx] = []
diff --git a/src/ednaml/backbones/resnet.py b/src/ednaml/backbones/resnet.py
index 82d0c851..d5490381 100644
--- a/src/ednaml/backbones/resnet.py
+++ b/src/ednaml/backbones/resnet.py
@@ -69,6 +69,7 @@ def __init__(
         self.secondary_attention = secondary_attention
         self.block = block
         self.inplanes = 64
+        self.initial_channels = kwargs.get("initial_channels", 3)
         if norm_layer is None:
             self._norm_layer = nn.BatchNorm2d
         # elif norm_layer == "ln":
@@ -85,7 +86,7 @@ def __init__(
         self.base_width = width_per_group
 
         self.conv1 = nn.Conv2d(
-            3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
+            self.initial_channels, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False
         )
         # if norm_layer == "gn":
         #    self.bn1 = nn.GroupNorm2d
@@ -279,6 +280,11 @@ def load_param(self, weights_path: str):
             weights_path (str): Path to the weights file
         """
         param_dict = torch.load(weights_path)
+        if self.initial_channels == 1:
+            conv1_weight = param_dict['conv1.weight']
+            param_dict['conv1.weight'] = conv1_weight.sum(dim=1, keepdim=True)
+        elif self.initial_channels != 3:
+            raise RuntimeError("Invalid number of input channels for pretrained weights. Need 1 or 3 channels, got %i"%self.initial_channels)
         for i in param_dict:
             if "fc" in i and self.top_only:
                 continue
diff --git a/src/ednaml/config/ConfigDefaults.py b/src/ednaml/config/ConfigDefaults.py
index b6a69b5d..b1c123b3 100644
--- a/src/ednaml/config/ConfigDefaults.py
+++ b/src/ednaml/config/ConfigDefaults.py
@@ -21,8 +21,9 @@ class ConfigDefaults(BaseConfig):
     CHECKPOINT_DIRECTORY: str
 
     OUTPUT_ARGS: Dict[str,str]
-    DEPLOYMENT: str
+    DEPLOY: str
     DEPLOYMENT_ARGS: Dict[str,str]
+    DEPLOYMENT_EPOCHS: int
 
     TRANSFORM_ARGS: Dict
     BATCH_SIZE: int
@@ -54,6 +55,14 @@ class ConfigDefaults(BaseConfig):
     LR_KWARGS: Dict[str, str]
 
     STEP_VERBOSE: int
+    INPUT_SIZE: List[int]
+    PLUGIN: str
+    PLUGIN_NAME: str
+    PLUGIN_KWARGS: Dict[str, str]
+
+    STORAGE_TYPE: str
+    STORAGE_ARGS: Dict
+    STORAGE_URL: str
 
     def __init__(self, **kwargs):
 
@@ -63,7 +72,7 @@ def __init__(self, **kwargs):
         self.MODEL_SERVING = kwargs.get("MODEL_SERVING", False)
         self.EPOCHS = kwargs.get("EPOCHS", 10)
         self.SKIPEVAL = kwargs.get("SKIPEVAL", False)
-        self.TEST_FREQUENCY = kwargs.get("TEST_FREQUENCY", 5)
+        self.TEST_FREQUENCY = kwargs.get("TEST_FREQUENCY", 1)
         self.TRAINER = kwargs.get("TRAINER", "BaseTrainer")
         self.TRAINER_ARGS = kwargs.get(
             "TRAINER_ARGS",
@@ -74,18 +83,19 @@ def __init__(self, **kwargs):
 
         self.MODEL_VERSION = kwargs.get("MODEL_VERSION", 1)
         self.MODEL_CORE_NAME = kwargs.get("MODEL_CORE_NAME", "model")
-        self.MODEL_BACKBONE = kwargs.get("MODEL_BACKBONE", "")
+        self.MODEL_BACKBONE = kwargs.get("MODEL_BACKBONE", "backbone")
         self.MODEL_QUALIFIER = kwargs.get("MODEL_QUALIFIER", "all")
         self.DRIVE_BACKUP = kwargs.get("DRIVE_BACKUP", False)
-        self.LOG_BACKUP = kwargs.get("LOG_BACKUP", True)
+        self.LOG_BACKUP = kwargs.get("LOG_BACKUP", False)
         self.SAVE_FREQUENCY = kwargs.get("SAVE_FREQUENCY", 5)
         self.CHECKPOINT_DIRECTORY = kwargs.get(
             "CHECKPOINT_DIRECTORY", "checkpoint"
         )
 
         self.OUTPUT_ARGS = {}
-        self.DEPLOYMENT = kwargs.get("DEPLOYMENT", "BaseDeploy")
+        self.DEPLOY = kwargs.get("DEPLOY", "BaseDeploy")
         self.DEPLOYMENT_ARGS = kwargs.get("DEPLOYMENT_ARGS", {})
+        self.DEPLOYMENT_EPOCHS = kwargs.get("DEPLOYMENT_EPOCHS", 1)
 
         self.TRANSFORM_ARGS = {}
         self.BATCH_SIZE = kwargs.get("BATCH_SIZE", 32)
@@ -118,3 +128,12 @@ def __init__(self, **kwargs):
         self.LR_KWARGS = kwargs.get("LR_KWARGS", {"step_size": 20})
 
         self.STEP_VERBOSE = kwargs.get("STEP_VERBOSE", 100)
+        self.INPUT_SIZE = kwargs.get("INPUT_SIZE", None)
+
+        self.STORAGE_TYPE = kwargs.get("STORAGE_TYPE", "BaseStorage")
+        self.STORAGE_ARGS = kwargs.get("STORAGE_ARGS", {})
+        self.STORAGE_URL = kwargs.get("STORAGE_URL", "./")
+        
+        self.PLUGIN_NAME = "mp-1"
+        self.PLUGIN = "ModelPlugin"
+        self.PLUGIN_KWARGS = {}
diff --git a/src/ednaml/config/DeploymentConfig.py b/src/ednaml/config/DeploymentConfig.py
index 625660eb..7300da94 100644
--- a/src/ednaml/config/DeploymentConfig.py
+++ b/src/ednaml/config/DeploymentConfig.py
@@ -1,6 +1,7 @@
 from ednaml.config import BaseConfig
 from ednaml.config.ConfigDefaults import ConfigDefaults
 from ednaml.config.ExecutionDatareaderConfig import ExecutionDatareaderConfig
+from ednaml.config.ExecutionPluginConfig import ExecutionPluginConfig
 from typing import Dict
 
 class DeploymentConfig(BaseConfig):
@@ -8,6 +9,9 @@ class DeploymentConfig(BaseConfig):
     DEPLOYMENT_ARGS: Dict[str,str]
     DEPLOY: str
     DATAREADER: ExecutionDatareaderConfig
+    EPOCHS: int
+    PLUGIN: ExecutionPluginConfig
+
 
     def __init__(self, deployment_dict, defaults: ConfigDefaults):
         self.OUTPUT_ARGS = deployment_dict.get(
@@ -16,9 +20,11 @@ def __init__(self, deployment_dict, defaults: ConfigDefaults):
         self.DEPLOYMENT_ARGS = deployment_dict.get(
             "DEPLOYMENT_ARGS", defaults.DEPLOYMENT_ARGS
         )
-        self.DEPLOYMENT = deployment_dict.get(
-            "DEPLOYMENT", defaults.DEPLOYMENT
+        self.DEPLOY = deployment_dict.get(
+            "DEPLOY", defaults.DEPLOY
         )
         self.DATAREADER = ExecutionDatareaderConfig(
             deployment_dict.get("DATAREADER", {})
-        )
\ No newline at end of file
+        )
+        self.EPOCHS = deployment_dict.get("EPOCHS", defaults.DEPLOYMENT_EPOCHS)
+        self.PLUGIN = ExecutionPluginConfig(deployment_dict.get("PLUGIN", {}))
\ No newline at end of file
diff --git a/src/ednaml/config/EdnaMLConfig.py b/src/ednaml/config/EdnaMLConfig.py
index c5a8d559..777b3799 100644
--- a/src/ednaml/config/EdnaMLConfig.py
+++ b/src/ednaml/config/EdnaMLConfig.py
@@ -1,10 +1,13 @@
 # BaseConfig class to manage the input configurations. TODO
 import os
 import json
+from turtle import update
 from typing import Dict, List
+import warnings
 import yaml
 from ednaml.config import BaseConfig
 from ednaml.config.DeploymentConfig import DeploymentConfig
+from ednaml.config.ModelPluginConfig import ModelPluginConfig
 from ednaml.utils import config_serializer
 from ednaml.config.ConfigDefaults import ConfigDefaults
 
@@ -16,6 +19,7 @@
 from ednaml.config.SchedulerConfig import SchedulerConfig
 from ednaml.config.TransformationConfig import TransformationConfig
 from ednaml.config.ModelConfig import ModelConfig
+from ednaml.config.StorageConfig import StorageConfig
 from ednaml.utils import merge_dictionary_on_key_with_copy
 
 class EdnaMLConfig(BaseConfig):
@@ -26,6 +30,7 @@ class EdnaMLConfig(BaseConfig):
     TEST_TRANSFORMATION: TransformationConfig
     MODEL: ModelConfig
     LOSS: List[LossConfig]
+    MODEL_PLUGIN: Dict[str, ModelPluginConfig]
     OPTIMIZER: List[
         OptimizerConfig
     ]  # one optimizer for each set of model params
@@ -37,62 +42,153 @@ class EdnaMLConfig(BaseConfig):
         SchedulerConfig
     ]  # one scheduler for each loss_optimizer
     LOGGING: LoggingConfig
+    STORAGE: StorageConfig
 
     extensions: List[str]
 
     def __init__(
-        self, config_path: str, defaults: ConfigDefaults = ConfigDefaults()
+        self, config_path: List[str], defaults: ConfigDefaults = ConfigDefaults(), **kwargs
     ):
-        self.extensions = ["DEPLOYMENT"]
-        ydict = self.read_path(config_path)
+        self.extensions = ["EXECUTION", "SAVE", "STORAGE", "TRANSFORMATION", "MODEL", "LOSS", "OPTIMIZER", "SCHEDULER", "LOSS_OPTIMIZER", "LOSS_SCHEDULER", "LOGGING", "DEPLOYMENT", "MODEL_PLUGIN"]  # TODO deal with other bits and pieces here!!!!!
+        ydict = self.merge([self.read_path(config_component) for config_component in config_path])
+        config_inject = kwargs.get("config_inject", None)
+        if config_inject is not None and type(config_inject) is list:
+            self.config_inject(ydict, config_inject)
+        self._updateConfig(ydict, defaults, update_with_defaults=True)
+    
+    def config_inject(self, ydict, config_inject: List[List[str]]):
+        for inject in config_inject:
+            inject_key = inject[0]
+            inject_value = inject[1]
+            if len(inject) == 3:
+                inject_kwarg = inject[2]
+            else:
+                inject_kwarg = None
+            try:
+                self._setinject(ydict, inject_key.split("."), inject_value, inject_kwarg)
+                print("Injected key-value pair:  {key}, {value}".format(key=inject_key, value=inject_value))
+            except KeyError:
+                pass
+
+
+    def _setinject(self, d, inject_key, inject_val, inject_kwarg = None):
+        for elem in inject_key[:-1]:
+            d = d[elem]
+        if inject_kwarg is not None:
+            d[inject_key[-1]][inject_kwarg] = inject_val    
+        else:
+            d[inject_key[-1]] = inject_val
+
 
-        self.EXECUTION = ExecutionConfig(ydict.get("EXECUTION", {}), defaults) #inside the exectution -- store execution object - -execution section
-        self.DEPLOYMENT = DeploymentConfig(ydict.get("DEPLOYMENT", {}), defaults) #inside the exectution -- store execution object - -execution section
-        self.SAVE = SaveConfig(ydict.get("SAVE", {}), defaults)
-        self.TRAIN_TRANSFORMATION = TransformationConfig(
-            dict(
-                merge_dictionary_on_key_with_copy(ydict.get("TRANSFORMATION", {}), ydict.get("TRAIN_TRANSFORMATION", {}))
-            ),
-            defaults,
-        )
-        #print("SELF.TRAIN_TRANSFORMATION ::::::::::::::::::: ",self.TRAIN_TRANSFORMATION)
-        self.TEST_TRANSFORMATION = TransformationConfig(
-            dict(
-                merge_dictionary_on_key_with_copy(ydict.get("TRANSFORMATION", {}), ydict.get("TEST_TRANSFORMATION", {}))
-            ),
-            defaults,
-        )
-        #print("SELF.TEST_TRANSFORMATION ::::::::::::::::::: ",self.TEST_TRANSFORMATION)
-
-        self.MODEL = ModelConfig(
-            ydict.get("MODEL", {}), defaults
-        )  # No default MODEL itself, though it will be instantiated here? deal with this TODO
-        self.LOSS = [
-            LossConfig(loss_item, defaults)
-            for loss_item in ydict.get("LOSS", [])
-        ]  # No default LOSS itself --> it will be empty...
-
-        # Default optimizer is Adam
-        self.OPTIMIZER = [
-            OptimizerConfig(optimizer_item, defaults)
-            for optimizer_item in ydict.get("OPTIMIZER", [{}])
-        ]
-        self.SCHEDULER = [
-            SchedulerConfig(scheduler_item, defaults)
-            for scheduler_item in ydict.get("SCHEDULER", [{}])
-        ]
-
-        self.LOSS_OPTIMIZER = [
-            OptimizerConfig(optimizer_item, defaults)
-            for optimizer_item in ydict.get("LOSS_OPTIMIZER", [{}])
-        ]
-        self.LOSS_SCHEDULER = [
-            SchedulerConfig(scheduler_item, defaults)
-            for scheduler_item in ydict.get("LOSS_SCHEDULER", [{}])
-        ]
-
-        self.LOGGING = LoggingConfig(ydict.get("LOGGING", {}), defaults)
+    def _has_extension_verifier(self, ydict, extension, verification):
+        return False if ydict.get(extension, verification) == verification else True
 
+    def _updateConfig(self, ydict: Dict[str,str], defaults: ConfigDefaults, update_with_defaults: bool):
+        added_extensions = []   # add control to check if extension was added or default was used...
+        for extension in self.extensions:
+            
+            if extension == "EXECUTION":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.EXECUTION = ExecutionConfig(ydict.get(extension, {}), defaults) 
+                    added_extensions.append([extension])
+            elif extension == "DEPLOYMENT":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.DEPLOYMENT = DeploymentConfig(ydict.get(extension, {}), defaults)
+                    added_extensions.append([extension])
+            elif extension == "STORAGE":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.STORAGE = StorageConfig(ydict.get(extension, {}), defaults)
+                    added_extensions.append([extension])
+            elif extension == "SAVE":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.SAVE = SaveConfig(ydict.get(extension, {}), defaults)
+                    added_extensions.append([extension])
+            elif extension == "TRANSFORMATION":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                has_train = self._has_extension_verifier(ydict, "TRAIN_TRANSFORMATION", {})
+                if (has_extension and has_train) or update_with_defaults:
+                    self.TRAIN_TRANSFORMATION = TransformationConfig(
+                        dict(
+                            merge_dictionary_on_key_with_copy(ydict.get(extension, {}), ydict.get("TRAIN_TRANSFORMATION", {}))
+                        ),
+                        defaults,
+                    )
+                    added_extensions.append(["TRAIN_TRANSFORMATION"])
+                has_test = self._has_extension_verifier(ydict, "TEST_TRANSFORMATION", {})
+                if (has_extension and has_test) or update_with_defaults:
+                    self.TEST_TRANSFORMATION = TransformationConfig(
+                        dict(
+                            merge_dictionary_on_key_with_copy(ydict.get(extension, {}), ydict.get("TEST_TRANSFORMATION", {}))
+                        ),
+                        defaults,
+                    )
+                    added_extensions.append(["TEST_TRANSFORMATION"])
+            elif extension == "MODEL":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.MODEL = ModelConfig(
+                        ydict.get(extension, {}), defaults
+                    )  # No default MODEL itself, though it will be instantiated here? deal with this TODO
+                    added_extensions.append([extension])
+            elif extension == "MODEL_PLUGIN":
+                has_extension = self._has_extension_verifier(ydict, extension, [])
+                if has_extension or update_with_defaults:
+                    mp_list = [ModelPluginConfig(plugin_item, defaults) for plugin_item in ydict.get(extension, [])]
+                    self.MODEL_PLUGIN = {item.PLUGIN_NAME: item for item in mp_list}
+                    added_extensions.append([extension])
+            elif extension == "LOSS":
+                has_extension = self._has_extension_verifier(ydict, extension, [])
+                if has_extension or update_with_defaults:
+                    self.LOSS = [
+                        LossConfig(loss_item, defaults)
+                        for loss_item in ydict.get(extension, [])
+                    ]  # No default LOSS itself --> it will be empty...
+                    added_extensions.append([extension])
+            elif extension == "OPTIMIZER":
+                # Default optimizer is Adam
+                has_extension = self._has_extension_verifier(ydict, extension, [{}])
+                if has_extension or update_with_defaults:
+                    self.OPTIMIZER = [
+                        OptimizerConfig(optimizer_item, defaults)
+                        for optimizer_item in ydict.get(extension, [{}])
+                    ]
+                    added_extensions.append([extension])
+            elif extension == "SCHEDULER":
+                has_extension = self._has_extension_verifier(ydict, extension, [{}])
+                if has_extension or update_with_defaults:
+                    self.SCHEDULER = [
+                        SchedulerConfig(scheduler_item, defaults)
+                        for scheduler_item in ydict.get(extension, [{}])
+                    ]
+                    added_extensions.append([extension])
+            elif extension == "LOSS_OPTIMIZER":
+                has_extension = self._has_extension_verifier(ydict, extension, [{}])
+                if has_extension or update_with_defaults:
+                    self.LOSS_OPTIMIZER = [
+                        OptimizerConfig(optimizer_item, defaults)
+                        for optimizer_item in ydict.get(extension, [{}])
+                    ]
+                    added_extensions.append([extension])
+            elif extension == "LOSS_SCHEDULER":
+                has_extension = self._has_extension_verifier(ydict, extension, [{}])
+                if has_extension or update_with_defaults:
+                    self.LOSS_SCHEDULER = [
+                        SchedulerConfig(scheduler_item, defaults)
+                        for scheduler_item in ydict.get(extension, [{}])
+                    ]
+                    added_extensions.append([extension])
+            elif extension == "LOGGING":
+                has_extension = self._has_extension_verifier(ydict, extension, {})
+                if has_extension or update_with_defaults:
+                    self.LOGGING = LoggingConfig(ydict.get(extension, {}), defaults)
+                    added_extensions.append([extension])
+            else:
+                raise RuntimeError("Somehow received unhandled extension %s"%str(extension))
+        return added_extensions
         # Things without defaults that MUST be provided: model ✅, train_dataloader, loss ✅, trainer TODO
 
     def export(self, mode="yaml"):
@@ -105,21 +201,30 @@ def export(self, mode="yaml"):
         elif mode == "dict":
             return dicta
 
-    def extend(self, config_path: str, defaults: ConfigDefaults = ConfigDefaults()):
+    def extend(self, config_path: str, defaults: ConfigDefaults = ConfigDefaults(), **kwargs):
         """Extends the existing config object with fields from the provided config
 
         Args:
             config_path (str): The path to the config file to extend with
         """        
-        
+        responses = {}
+        if type(config_path) is str:
+            config_path = [config_path]
+        if type(config_path) is list:
+            for cpath in config_path:
+                responses[cpath] = self._extend(cpath, defaults, **kwargs)
+        else:
+            raise ValueError("Expected `list`, got %s"%type(config_path))
+
+        return responses
+
+    def _extend(self, config_path, defaults: ConfigDefaults, **kwargs):
         ydict = self.read_path(config_path)
-        for extension in self.extensions:
-            if len(ydict.get(extension, {})) > 0:
-                if extension == "DEPLOYMENT":
-                    self.DEPLOYMENT = DeploymentConfig(ydict.get(extension, {}), defaults) 
-                    return "Extended with DEPLOYMENT."
-                else:
-                    return "No valid extensions available."
+        config_inject = kwargs.get("config_inject", None)
+        if config_inject is not None and type(config_inject) is list:
+            self.config_inject(ydict, config_inject)
+        added_extensions = self._updateConfig(ydict, defaults, update_with_defaults=False)
+        return "Extended with : %s"%", ".join([item[0] for item in added_extensions if len(item)>0])
 
 
 
@@ -147,6 +252,63 @@ def read_path(self, path: os.PathLike) -> Dict[str,str]:
                     ydict = yaml.safe_load(cfile.read().strip()) #loading the configuration file.
         return ydict
 
+    def _merge(self,base_config: Dict[str,str], extension: Dict[str,str], recursion_flag: bool = False):
+        """Combines two config files
+        If capital lettered key is present it'll be replaced
+        If small lettered key is present, it'll be deleted and considered the key in second object only
+
+        Args:
+            first object: Dict (base config)
+            second object: Dict (second config)
+        Returns:
+            Dict[str,str]: Config as a combined python dictionary
+        """   
+        #if(recursion_flag):
+        #    to_keep = lambda key: key.islower()
+        #    {k: v for k, v in base_config.items() if k not in keyfilter(to_keep, base_config)}
+        lower_flag = False
+        upper_flag = False
+        for extension_key in extension.keys():
+            if extension_key in base_config:
+                if extension_key.islower(): # If NOT a built-in key, we will do a full replacement of the parent...
+                    lower_flag = True
+                else:   # If it is a built-in key
+                    upper_flag = True
+                    if isinstance(base_config[extension_key], dict):    # If it is a dict
+                        if isinstance(extension[extension_key], dict):
+                            # Recurse if key is a Dictionary with elements inside.
+                            base_config[extension_key] = self._merge(base_config[extension_key], extension[extension_key],recursion_flag=True)
+                        else:
+                            base_config[extension_key]=extension[extension_key]
+                            warnings.warn("Key {key}: Dict key in base configuration is not a dictionary in replacement extension. \n\tBase: {base_str}\n\tExtension: {ext_str}".format(
+                                key=extension_key,
+                                base_str = str(base_config(extension_key)),
+                                ext_str = str(base_config(extension_key)),
+                            ))
+                    else:   # The base config key is not a dictionary...we just do replacement. Don't care if replacement is dictionary
+                        base_config[extension_key]=extension[extension_key] 
+            else:   # a key in extensionis NOT in base_config, so we don't really care...
+                base_config[extension_key] = extension[extension_key]
+        if lower_flag and upper_flag:
+            raise ValueError("built-in and custom keys cannot mix")
+        if lower_flag:
+            return extension
+        return base_config
+
+    def merge(self, config_paths: List[str]):
+        """Combines multiple config files
+        Args:
+            first object: Dict array (base config as index 0)
+        Returns:
+            Dict[str,str]: Config as a combined python dictionary
+        """  
+        base_config = config_paths[0]
+        for extension_config in config_paths[1:]:
+            base_config = self._merge(base_config, extension_config)
+        return base_config
+
+
+
 """
 Notes for LOSS_OPTIMIZER and LOSS_SCHEDULER
         Now we import the loss-optimizers and loss-scheduler
diff --git a/src/ednaml/config/ExecutionConfig.py b/src/ednaml/config/ExecutionConfig.py
index fa159624..779381ef 100644
--- a/src/ednaml/config/ExecutionConfig.py
+++ b/src/ednaml/config/ExecutionConfig.py
@@ -1,7 +1,7 @@
 from ednaml.config import BaseConfig
 from ednaml.config.ConfigDefaults import ConfigDefaults
 from ednaml.config.ExecutionDatareaderConfig import ExecutionDatareaderConfig
-
+from ednaml.config.ExecutionPluginConfig import ExecutionPluginConfig
 
 class ExecutionConfig(BaseConfig):
     OPTIMIZER_BUILDER: str
@@ -11,6 +11,7 @@ class ExecutionConfig(BaseConfig):
     TEST_FREQUENCY: int
     TRAINER: str
     DATAREADER: ExecutionDatareaderConfig
+    PLUGIN: ExecutionPluginConfig
 
     def __init__(self, execution_dict, defaults: ConfigDefaults):
         self.OPTIMIZER_BUILDER = execution_dict.get(
@@ -33,3 +34,4 @@ def __init__(self, execution_dict, defaults: ConfigDefaults):
         self.TRAINER_ARGS = execution_dict.get(
             "TRAINER_ARGS", defaults.TRAINER_ARGS
         )
+        self.PLUGIN = ExecutionPluginConfig(execution_dict.get("PLUGIN", {}))
diff --git a/src/ednaml/config/ExecutionDatareaderConfig.py b/src/ednaml/config/ExecutionDatareaderConfig.py
index 2a7bcc9b..378b4dd9 100644
--- a/src/ednaml/config/ExecutionDatareaderConfig.py
+++ b/src/ednaml/config/ExecutionDatareaderConfig.py
@@ -9,7 +9,7 @@ class ExecutionDatareaderConfig(BaseConfig):
     DATASET_ARGS: Dict[str, str]
     GENERATOR: str
     GENERATOR_ARGS: Dict[str, str]
-    DATALOADER_ARGS: Dict[str, str]
+    
 
     def __init__(self, datareader_dict):
         self.DATAREADER = datareader_dict.get("DATAREADER", "DataReader")
@@ -19,4 +19,3 @@ def __init__(self, datareader_dict):
             "GENERATOR", None
         )
         self.GENERATOR_ARGS = datareader_dict.get("GENERATOR_ARGS", {})
-        self.DATALOADER_ARGS = datareader_dict.get("DATALOADER_ARGS", {})
diff --git a/src/ednaml/config/ExecutionPluginConfig.py b/src/ednaml/config/ExecutionPluginConfig.py
new file mode 100644
index 00000000..3943749f
--- /dev/null
+++ b/src/ednaml/config/ExecutionPluginConfig.py
@@ -0,0 +1,14 @@
+from typing import Dict
+
+from ednaml.config import BaseConfig
+
+
+class ExecutionPluginConfig(BaseConfig):
+    #EPOCHS: int
+    HOOKS: str   # always | warmup | activated
+    RESET: bool
+
+    def __init__(self, ep_dict):
+        #self.EPOCHS = datareader_dict.get("PLUGIN_EPOCHS", 1)
+        self.HOOKS = ep_dict.get("HOOKS", 'always')
+        self.RESET = ep_dict.get("RESET", False) # TODO
\ No newline at end of file
diff --git a/src/ednaml/config/LoggingConfig.py b/src/ednaml/config/LoggingConfig.py
index 90651557..97034b1d 100644
--- a/src/ednaml/config/LoggingConfig.py
+++ b/src/ednaml/config/LoggingConfig.py
@@ -9,3 +9,6 @@ def __init__(self, logging_dict, defaults: ConfigDefaults):
         self.STEP_VERBOSE = logging_dict.get(
             "STEP_VERBOSE", defaults.STEP_VERBOSE
         )
+        self.INPUT_SIZE = logging_dict.get(
+            "INPUT_SIZE", defaults.INPUT_SIZE
+        )
diff --git a/src/ednaml/config/ModelPluginConfig.py b/src/ednaml/config/ModelPluginConfig.py
new file mode 100644
index 00000000..e192342a
--- /dev/null
+++ b/src/ednaml/config/ModelPluginConfig.py
@@ -0,0 +1,22 @@
+from typing import Dict
+from ednaml.config import BaseConfig
+from ednaml.config.ConfigDefaults import ConfigDefaults
+
+
+class ModelPluginConfig(BaseConfig):
+    PLUGIN: str
+    PLUGIN_NAME: str
+    PLUGIN_KWARGS: Dict[str, str]
+    
+    
+    def __init__(self, save_dict, defaults: ConfigDefaults):
+        self.PLUGIN = save_dict.get(
+            "PLUGIN", defaults.PLUGIN
+        )
+        self.PLUGIN_NAME = save_dict.get(
+            "PLUGIN_NAME", defaults.PLUGIN_NAME
+        )
+        self.PLUGIN_KWARGS = save_dict.get(
+            "PLUGIN_KWARGS", defaults.PLUGIN_KWARGS
+        )
+        
\ No newline at end of file
diff --git a/src/ednaml/config/StorageConfig.py b/src/ednaml/config/StorageConfig.py
new file mode 100644
index 00000000..2dbc22ca
--- /dev/null
+++ b/src/ednaml/config/StorageConfig.py
@@ -0,0 +1,17 @@
+from ednaml.config import BaseConfig
+from typing import Dict
+
+from ednaml.config.ConfigDefaults import ConfigDefaults
+
+
+class StorageConfig(BaseConfig):
+    TYPE: str
+    STORAGE_ARGS: Dict
+    URL: str
+
+    def __init__(self, storage_dict,defaults: ConfigDefaults):
+        self.TYPE = storage_dict.get("TYPE", defaults.STORAGE_TYPE)
+        self.STORAGE_ARGS = storage_dict.get("STORAGE_ARGS", defaults.STORAGE_ARGS)
+        self.URL = storage_dict.get("URL", defaults.STORAGE_URL)
+        
+
diff --git a/src/ednaml/config/__init__.py b/src/ednaml/config/__init__.py
index 62f732ca..61bf641b 100644
--- a/src/ednaml/config/__init__.py
+++ b/src/ednaml/config/__init__.py
@@ -10,7 +10,7 @@ def __eq__(self, other):
         return self.__dict__ == other.__dict__
 
     def __repr__(self):
-        return str(self.to_json_string())
+        return self.to_json_string()
 
     def to_dict(self):
         """Serializes this instance to a Python dictionary."""
@@ -19,7 +19,7 @@ def to_dict(self):
 
     def to_json_string(self):
         """Serializes this instance to a JSON string."""
-        return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
+        return json.dumps(self.to_dict(), indent=2, sort_keys=True, default=lambda o: o.__dict__)
 
     def to_json_file(self, json_file_path):
         """Save this instance to a json file."""
diff --git a/src/ednaml/core/EdnaDeploy.py b/src/ednaml/core/EdnaDeploy.py
index 7f341f0a..c5a360fa 100644
--- a/src/ednaml/core/EdnaDeploy.py
+++ b/src/ednaml/core/EdnaDeploy.py
@@ -1,4 +1,4 @@
-from typing import Type
+from typing import Type, Union, List
 from ednaml.core import EdnaML
 import logging
 from ednaml.crawlers import Crawler
@@ -12,18 +12,19 @@ class EdnaDeploy(EdnaML):
     def __init__(
         self,
         config: str = "",
-        deploy: str = "",
+        deploy: Union[List[str], str] = "",
         mode: str = "test",
         weights: str = None,
         logger: logging.Logger = None,
         verbose: int = 2,
         **kwargs
     ):
-
+        
         super().__init__([config,deploy], mode, weights, logger, verbose,**kwargs)
 
-    
-
+        self.decorator_reference["deployment"] = self.addDeploymentClass
+        self.decorator_reference.pop("trainer") # We do not need trainer in a Deployment
+        self.dataloader_mode = kwargs.get("dataloader_mode", self.mode)
 
     def apply(self, **kwargs):
         """Applies the internal configuration for EdnaDeploy"""
@@ -104,11 +105,12 @@ def buildTestDataloader(
         if self._testGeneratorInstanceQueueFlag:
             self.test_generator: Generator = self._testGeneratorInstanceQueue
         else:
-            self.test_generator: Generator = data_reader.GENERATOR(
+            self.logger.info("Generating dataloader `{dataloader}` with `{mode}` mode".format(mode=self.dataloader_mode, dataloader=data_reader.GENERATOR.__name__))
+            self.test_generator: Type[Generator] = data_reader.GENERATOR(
                 logger=self.logger,
                 gpus=self.gpus,
                 transforms=self.cfg.TEST_TRANSFORMATION,
-                mode="test",
+                mode=self.dataloader_mode,  # TODO convert this to better options: i.e. which mode to use, and which transformations to use, as an option in data_reader, specifically for deployments
                 **self.cfg.DEPLOYMENT.DATAREADER.GENERATOR_ARGS
             )
             self.test_generator.build( 
@@ -135,13 +137,13 @@ def _buildCrawlerInstance(self, data_reader: DataReader) -> Crawler:
             logger=self.logger, **self.cfg.DEPLOYMENT.DATAREADER.CRAWLER_ARGS
         )
 
-    def deploy(self):
-        self.deployment.deploy(continue_epoch=self.previous_stop + 1)
+    def deploy(self, **kwargs):
+        self.deployment.deploy(continue_epoch=self.previous_stop + 1, **kwargs)
 
     def buildDeployment(self):
         """Builds the EdnaDeploy deployment and sets it up"""
         if self._deploymentClassQueueFlag:
-            ExecutionDeployment = self._deploymentClassQueueFlag
+            ExecutionDeployment = self._deploymentClassQueue
         else:
             ExecutionDeployment: Type[BaseDeploy] = locate_class(
                 subpackage="deploy", classpackage=self.cfg.DEPLOYMENT.DEPLOY
@@ -158,6 +160,7 @@ def buildDeployment(self):
             self.deployment = ExecutionDeployment(
                 model=self.model,
                 data_loader=self.test_generator.dataloader, # TODO
+                epochs=self.cfg.DEPLOYMENT.EPOCHS,
                 logger=self.logger,
                 crawler=self.crawler,
                 config=self.cfg,
@@ -168,6 +171,7 @@ def buildDeployment(self):
                 step_verbose=self.cfg.LOGGING.STEP_VERBOSE,
                 save_directory=self.saveMetadata.MODEL_SAVE_FOLDER,
                 save_backup=self.cfg.SAVE.DRIVE_BACKUP,
+                save_frequency=self.cfg.SAVE.SAVE_FREQUENCY,
                 backup_directory=self.saveMetadata.CHECKPOINT_DIRECTORY,
                 gpus=self.gpus,
                 fp16=self.cfg.EXECUTION.FP16,
@@ -186,6 +190,7 @@ def addResetQueues(self):
         return [self.resetDeploymentQueue]
 
     def addDeploymentClass(self, deployment_class: Type[BaseDeploy]):
+        self.logger.debug("Added deployment class: %s"%deployment_class.__name__)
         self._deploymentClassQueue = deployment_class
         self._deploymentClassQueueFlag = True
 
diff --git a/src/ednaml/core/EdnaML.py b/src/ednaml/core/EdnaML.py
index ecf1f555..7b20b72d 100644
--- a/src/ednaml/core/EdnaML.py
+++ b/src/ednaml/core/EdnaML.py
@@ -1,7 +1,8 @@
 import importlib
 import os, shutil, logging, glob, re
-from types import FunctionType, MethodType
+from types import MethodType
 from typing import Callable, Dict, List, Type, Union
+from urllib.error import URLError
 import warnings
 from torchinfo import ModelStatistics
 from ednaml.config.EdnaMLConfig import EdnaMLConfig
@@ -16,8 +17,10 @@
 from ednaml.optimizer import BaseOptimizer
 from ednaml.optimizer.StandardLossOptimizer import StandardLossOptimizer
 from ednaml.loss.builders import ClassificationLossBuilder
+from ednaml.plugins.ModelPlugin import ModelPlugin
 from ednaml.trainer.BaseTrainer import BaseTrainer
-from ednaml.utils import locate_class
+from ednaml.storage import BaseStorage
+from ednaml.utils import locate_class, path_import
 import ednaml.utils
 import torch
 from torchinfo import summary
@@ -25,6 +28,7 @@
 import ednaml.utils.web
 from ednaml.utils.SaveMetadata import SaveMetadata
 import logging
+import ednaml.core.decorators
 
 """TODO
 
@@ -35,7 +39,7 @@
 class EdnaML(EdnaMLBase):
     labelMetadata: LabelMetadata
     modelStatistics: ModelStatistics
-    model: ModelAbstract
+    model: ModelAbstract = None
     loss_function_array: List[LossBuilder]
     loss_optimizer_array: List[torch.optim.Optimizer]
     optimizer: List[torch.optim.Optimizer]
@@ -47,6 +51,8 @@ class EdnaML(EdnaMLBase):
     train_generator: Generator
     test_generator: Generator
     cfg: EdnaMLConfig
+    decorator_reference: Dict[str,Type[MethodType]]
+    plugins: Dict[str, ModelPlugin] = {}
 
     def __init__(
         self,
@@ -76,22 +82,31 @@ def __init__(
 
         """
 
-        self.config = config
+        if type(config) is str:
+            self.config = config
+        elif type(config) is list:
+            self.config = []
+            for cpath in config:
+                if type(cpath) is str:
+                    self.config.append(cpath)
+                elif type(cpath) is list:
+                    self.config += cpath
+                else:
+                    ValueError("config MUST be list or string")
+        else:
+            raise ValueError("config MUST be list or string")
+        
         self.mode = mode
         self.weights = weights
         self.pretrained_weights = None
         self.verbose = verbose
         self.gpus = torch.cuda.device_count()
-
-        # Added configuration extentions
-        if type(self.config) is list:
-            self.cfg = EdnaMLConfig(config[0])
-            msg = self.cfg.extend(config[1])
-            print(msg)
+        # TODO Deal with extensions
+        if type(self.config) is str:
+            self.cfg = EdnaMLConfig([self.config], **kwargs)
         else:
-            self.cfg = EdnaMLConfig(config) 
-        
-        
+            self.cfg = EdnaMLConfig(self.config, **kwargs)
+
         self.saveMetadata = SaveMetadata(
             self.cfg, **kwargs
         )  # <-- for changing the logger name...
@@ -106,10 +121,19 @@ def __init__(
                 "Cannot have `test_only` and be in training mode. Switch to"
                 " `test` mode."
             )
-        self.resetQueueArray:List[MethodType] = [self.resetCrawlerQueues, self.resetGeneratorQueues, self.resetModelQueues, self.resetOptimizerQueues, self.resetLossBuilderQueue, self.resetTrainerQueue]
+        self.resetQueueArray:List[MethodType] = [self.resetCrawlerQueues, self.resetGeneratorQueues, self.resetModelQueues, self.resetOptimizerQueues, self.resetLossBuilderQueue, self.resetStorageQueues, self.resetTrainerQueue]
         self.resetQueueArray += self.addResetQueues()
         self.resetQueues()
 
+        self.decorator_reference = {
+            "crawler": self.addCrawlerClass,
+            "model": self.addModelClass,
+            "trainer": self.addTrainerClass,
+            #"storage": self.addStorageClass,
+            "generator": self.addGeneratorClass,
+            "model_plugin": self.addPlugins,
+        }
+
     def addResetQueues(self):
         return []
     def resetCrawlerQueues(self):
@@ -151,9 +175,16 @@ def resetTrainerQueue(self):
         self._trainerClassQueueFlag = False
         self._trainerInstanceQueue = None
         self._trainerInstanceQueueFlag = False
+    def resetStorageQueues(self):
+        self._storageInstanceQueueFlag = False
+        self._storageClassQueueFlag = False
+        self._storageInstanceQueue = None
+        self._storageClassQueue = None
+
 
     def resetQueues(self):
         """Resets the `apply()` queue"""
+        self.logger.debug("Resetting declarative queues.")
         for queue_function in self.resetQueueArray:
             queue_function()
 
@@ -223,12 +254,17 @@ def apply(self, **kwargs):
         self.printConfiguration()
         self.downloadModelWeights()
         self.setPreviousStop()
+        self.buildStorage() # this is where -- same as other build things...  
 
         self.buildDataloaders()
 
         self.buildModel()
         self.loadWeights()
-        self.getModelSummary(**kwargs) 
+        if not kwargs.get("skip_model_summary", False):
+            if self.cfg.LOGGING.INPUT_SIZE is not None:
+                if kwargs.get("input_size", None) is None:
+                    kwargs["input_size"] = self.cfg.LOGGING.INPUT_SIZE
+            self.getModelSummary(**kwargs) 
         self.buildOptimizer()
         self.buildScheduler()
 
@@ -236,17 +272,57 @@ def apply(self, **kwargs):
         self.buildLossOptimizer()
         self.buildLossScheduler()
 
+        self.buildStorage()
         self.buildTrainer()
 
         self.resetQueues()
 
-    def train(self):
-        self.trainer.train(continue_epoch=self.previous_stop + 1)  #
+    def addStorage(self, storage: BaseStorage):
+        self.logger.debug("Added storage: %s"%storage.__class__.__name__)
+        self._storageInstanceQueue = storage
+        self._storageInstanceQueueFlag = True
+    
+    def addStorageClass(self, storage_class: Type[BaseStorage]):
+        self.logger.debug("Added storage class: %s"%storage_class.__name__)
+        self._storageClassQueue = storage_class
+        self._storageClassQueueFlag = True
+
+    def buildStorage(self):  
+        # Options: a class was passed through the edna add functionality
+        # Alternatively, we have specified a specific class in the storage.type section
+        # Alternatively, we have specified a type, but that does bot exist as a built in
+
+        # So, first we check if a class or instance was directly passed.
+        # If nothing has been passed, then we try to locate a built-in version based on storage.type
+        # If that doesn't work, then we can throw an error...
+
+        if self._storageClassQueueFlag:
+            storage = self._storageClassQueue
+        else:
+            storage: Type[BaseStorage] = locate_class(
+                subpackage="storage", classpackage=self.cfg.STORAGE.TYPE
+            )
+            self.logger.info(
+                "Loaded {} from {} to build Storage".format(
+                    self.cfg.STORAGE.TYPE, "ednaml.storage"
+                )
+            )
+
+        if self._storageInstanceQueueFlag:
+            self.storage = self._storageInstanceQueue
+        else:
+            self.storage = storage(type=self.cfg.STORAGE.TYPE,
+                                    url=self.cfg.STORAGE.URL,
+                                    **self.cfg.STORAGE.STORAGE_ARGS)
+
+    def train(self, **kwargs):
+        self.trainer.train(continue_epoch=self.previous_stop + 1, **kwargs)  #
 
     def eval(self):
         return self.trainer.evaluate()
 
     def addTrainerClass(self, trainerClass: Type[BaseTrainer]):
+        self.logger.debug("Added trainer class: %s"%trainerClass.__name__)
         self._trainerClassQueue = trainerClass
         self._trainerClassQueueFlag = True
 
@@ -286,9 +362,10 @@ def buildTrainer(self):
                 crawler=self.crawler,
                 config=self.cfg,
                 labels=self.labelMetadata,
+                storage=self.storage,
                 **self.cfg.EXECUTION.TRAINER_ARGS
             )
-            self.trainer.setup(
+            self.trainer.apply(
                 step_verbose=self.cfg.LOGGING.STEP_VERBOSE,
                 save_frequency=self.cfg.SAVE.SAVE_FREQUENCY,
                 test_frequency=self.cfg.EXECUTION.TEST_FREQUENCY,
@@ -648,7 +725,60 @@ def buildModel(self):
                 self.cfg.MODEL.MODEL_ARCH
             )
         )
+        self.model._logger = self.logger
+        self.logger.info("Adding plugins after constructing model")
+        self.addPluginsToModel()
+
+    # Add model hooks here, since it is a ModelAbstract
+    def addPlugins(self, plugin_class_list: List[Type[ModelPlugin]], **kwargs):
+        """add plugins to the plugins class queue...
+        ideally this is called before add model
+        then add model can add the plugins in the queue to the model...
+
+        NOTE -- any plugins are added through the config file, specifically, with the plugin arguments. 
+        This is because multiple plugins can use the same class, but variations of it.
+        For example, one might want to train KMP under l2 and cos, and use then simultaneously...
+
+        So, plugins class queue exists only as a lookup for plugin classes...if a correspondiing enttry is NOT in MODEL_PLUGINS in the config,
+        it will never be added into Edna/Deploy...
+        """
+        for plugin in plugin_class_list:
+            self.logger.debug("Added custom plugin: %s"%plugin.__name__)
+            self.plugins[plugin.__name__] = plugin  # order matters; can replace...though shouldn't matter too much...
+        
+        # The follow will not happen, logically, i think, because model is ONLY defined in self.buildModel(), and by that point, one is already calling apply()
+        # Adding plugins after the fact...????? Don't think so...
+        # And in any case, we should make apply() truely declarative, so that one can simply apply() again after adding plugins
+        #if self.model is not None:
+        #    # model is defined. We will add plugins to queue, and then call the model's add plugin bit...
+        #    self.logger.info("Model already constructed. Adding plugins.")
+        #    self.addPluginsToModel()
+
+    def addPluginsToModel(self):
+        # Here, we iterate through plugins IN the config file, then import their classes, and add them to the model...
+        for plugin_save_name in self.cfg.MODEL_PLUGIN:
+            # Now, we need to locate the plugin class...check first if it is in plugins...
+            if self.cfg.MODEL_PLUGIN[plugin_save_name].PLUGIN not in self.plugins:
+                # now we need to locate the class in builtins
+                plugin = locate_class(
+                    subpackage="plugins",
+                    classpackage=self.cfg.MODEL_PLUGIN[plugin_save_name].PLUGIN,
+                    classfile=self.cfg.MODEL_PLUGIN[plugin_save_name].PLUGIN
+                )
+            else:
+                plugin = self.plugins[self.cfg.MODEL_PLUGIN[plugin_save_name].PLUGIN]
+            
+            self.model.addPlugin(
+                plugin = plugin,
+                plugin_name = plugin_save_name,
+                plugin_kwargs = self.cfg.MODEL_PLUGIN[plugin_save_name].PLUGIN_KWARGS
+            )
 
+        #for plugin in self.plugins:
+        #    self.model.addPlugin(plugin, 
+        #                    plugin_name = self.cfg.MODEL_PLUGIN[plugin.name].PLUGIN_NAME, 
+        #                    plugin_kwargs = self.cfg.MODEL_PLUGIN[plugin.name].PLUGIN_KWARGS)
+            
     def addModelBuilder(
         self, model_builder: Type[Callable], model_config: ModelConfig = None
     ):
@@ -663,6 +793,7 @@ def addModel(self, model: ModelAbstract):
         self._modelQueueFlag = True
 
     def addModelClass(self, model_class: Type[ModelAbstract], **kwargs):
+        self.logger.debug("Added model class: %s"%model_class.__name__)
         self._modelClassQueue = model_class
         self._modelClassQueueFlag = True
         self._modelArgsQueue = kwargs
@@ -731,32 +862,80 @@ def getModelSummary(self, input_size=None, dtypes=None, **kwargs):
         # add bookkeeping
         # does config has input size? if it does use that, but prepend batch size .
         # if it doesn't have it, use input size in arguments
-        self.model.cuda()
+        # TODO possibly move this into trainer...?
+        # or at least deal with potential mlti-gpu scenario...
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.model.to(device)
         # change below statement according to line 722
         # default for input size is None/null
-        if input_size is None:
-            input_size = (
-                self.cfg.TRAIN_TRANSFORMATION.BATCH_SIZE,
-                self.cfg.TRAIN_TRANSFORMATION.INPUT_SIZE,
-            ) # INPUT SIZE SHOULD HAVE A VALUE
-        print("INPUT SIZE ==== ",input_size)
-        self.model_summary = summary(
-            self.model,
-            input_size=input_size,
-            col_names=[
-                "input_size",
-                "output_size",
-                "num_params",
-                "kernel_size",
-                "mult_adds",
-            ],
-            depth=3,
-            mode="train",
-            verbose=0,
-            dtypes=dtypes,
-        )
-        self.logger.info(str(self.model_summary))
-
+        try:
+            if input_size is None:
+                input_size = (
+                    self.cfg.TRAIN_TRANSFORMATION.BATCH_SIZE,
+                    self.cfg.TRAIN_TRANSFORMATION.INPUT_SIZE,
+                ) # INPUT SIZE SHOULD HAVE A VALUE
+            print("INPUT SIZE ==== ",input_size)
+
+            self.model_summary = summary(
+                self.model,
+                input_size=input_size,
+                col_names=[
+                    "input_size",
+                    "output_size",
+                    "num_params",
+                    "kernel_size",
+                    "mult_adds",
+                ],
+                depth=3,
+                mode="train",
+                verbose=0,
+                dtypes=dtypes,
+            )
+            self.logger.info(str(self.model_summary))
+        except KeyboardInterrupt:   # TODO check which exception is actually raised in summary and maybe have a better way to deal with this...
+            raise
+        except Exception as e:
+            import traceback
+            self.logger.info("Model Summary retured the following error:")
+            self.logger.info(traceback.format_exc())
+
+
+    #------------------------------GENERAL ADD--------------------------------------------------------------------
+    def _add(self, file_or_module_path):
+        # We are given a file, which contains several class loaded through decorators when the file is imported.
+        imported_file = path_import(file_or_module_path)
+        # Once we have imported, the files are registered in ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS
+        lookup_path = os.path.abspath(file_or_module_path)
+        for keyvalue in ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path]:
+            if keyvalue in self.decorator_reference:
+                if type(ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue]) is list:
+                  for i in range(len(ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue])):
+                    self.logger.info("Adding a {ftype}, from {src}, with inferred name {inf}".format(
+                    ftype=keyvalue,
+                    src = lookup_path,
+                    inf = ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue][i].__name__
+                ))
+                else:
+                  self.logger.info("Adding a {ftype}, from {src}, with inferred name {inf}".format(
+                      ftype=keyvalue,
+                      src = lookup_path,
+                      inf = ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue].__name__
+                  ))
+                self.decorator_reference[keyvalue](
+                    ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue]
+                )
+            else:
+                warnings.warn(
+                    "keyvalue %s in REGISTERED_EDNA_COMPONENTS %s is not available in self.decorator_reference. Not adding."%(keyvalue, 
+                    str(ednaml.core.decorators.REGISTERED_EDNA_COMPONENTS[lookup_path][keyvalue]))
+                )
+    def add(self, file_or_module_path):
+        if type(file_or_module_path) is list:
+            for file_or_module in file_or_module_path:
+                self._add(file_or_module)
+        else:
+            self._add(file_or_module_path)
+        
     # ----------------------------------------------   DATAREADERS   ----------------------------------------------
     def addCrawlerClass(self, crawler_class: Type[Crawler], **kwargs):
         """Adds a crawler class to the EdnaML `apply()` queue. This will be applied to the configuration when calling `apply()`
@@ -892,14 +1071,7 @@ def buildTrainDataloader(
         if self._trainGeneratorInstanceQueueFlag:
             self.train_generator: Generator = self._trainGeneratorInstanceQueue
         else:
-            #print("??????????????????????????????????????????????????????????????????????//",self.cfg.TRAIN_TRANSFORMATION.getVars())
-            #train_transform_values = self.cfg.TRAIN_TRANSFORMATION.getVars()
-            #values_for_transform = {}
-            #values_for_transform['i_shape'] = train_transform_values['SHAPE']
-            #values_for_transform['i_shape'] = train_transform_values['SHAPE']
-            #values_for_transform['i_shape'] = train_transform_values['SHAPE']
-            #print("THIS IS A CHANGE!")
-            print(self.cfg.TRAIN_TRANSFORMATION)
+            #print(self.cfg.TRAIN_TRANSFORMATION)
             if self.mode != "test":
                 self.train_generator: Generator = data_reader.GENERATOR( ##imp -- initialize generator
                     logger=self.logger,
@@ -916,10 +1088,13 @@ def buildTrainDataloader(
                     **self.cfg.EXECUTION.DATAREADER.DATASET_ARGS
                 )
             else:
-                self.train_generator = Generator()
+                self.logger.info(
+                "Not creating `train_generator` in `test_only` mode."
+                )
+                self.train_generator = Generator(logger=self.logger)
         if self.mode != "test":
             self.logger.info(
-                "Generated training data generator with %i trainnig data points"
+                "Generated training data generator with %i training data points"
                 % len(self.train_generator.dataset)
             )
             self.labelMetadata = self.train_generator.num_entities
@@ -1121,4 +1296,4 @@ def deleteLocal(self):
         shutil.rmtree(self.saveMetadata.MODEL_SAVE_FOLDER)
 
     def setBackupToFalse(self):
-        self.cfg.SAVE.DRIVE_BACKUP = False
+        self.cfg.SAVE.DRIVE_BACKUP = False
\ No newline at end of file
diff --git a/src/ednaml/core/ModelConversionDeployment.py b/src/ednaml/core/ModelConversionDeployment.py
new file mode 100644
index 00000000..04cd2731
--- /dev/null
+++ b/src/ednaml/core/ModelConversionDeployment.py
@@ -0,0 +1,53 @@
+"""
+ModelConversionDeployment is a built-in Edna class to convert models between different ModelAbstract implementations.
+
+The use case for such a ModelConversionDeployment is to change how a model functions after it has already been trained, such as adding additional outputs like moden confidence or self-diagnoses, adjusting how a model manipulates its input, including a post-facto reject option, etc.
+
+Prerequisites for using ModelConversionDeployment:
+
+- a trained model's saved weights
+
+Corequisites for using ModelConversionDeployment:
+
+- a model class, inherited from ModelAbstract
+    - ModelAbstract already contains methods to load weights from a .pth file where the keys do not exactly match, with partial_load. This would normally be used to load the saved weights of the older model into the new model
+- a config file that links to the old weights through the Edna section OR the save section
+- NOTE --> we still need to deal with one small tiny issue!!!!!!! when providing multiple configs, the latter replace the older on same keys. So will the deployment.yml's SAVE section replace the model.yml's SAVE section????
+- a config file with an updated Save section to push the new model and weights
+- any helper files -- if not provided, ModelConversionDeployment will use the classes from the original models' implementation, unless a new Crawler is provided
+- 
+
+
+Usage Instructions:
+
+
+edna ModelConversionDeployment new_model_file.py model_config.yml deployment_config.yml
+
+"""
+
+from ednaml.deploy import BaseDeploy
+# So we need the load() to be partial load, possibly...
+# WE NEED TO ADJUST HOW MODEL LOADS ARE HANDLED IN EDNAML AND EDNADEPLOY
+
+class ModelConversionDeployment(BaseDeploy):
+  def deploy_step(self, batch):
+    batch = tuple(item.cuda() for item in batch)
+    all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_labels = batch
+    # outputs is a tuple of output, features, secondary_output --> which has the plugin components. Maybe we will split plugin outputs separately...
+    outputs = self.model(all_input_ids, token_type_ids = all_token_type_ids, attention_mask=all_attention_mask)
+    logit = outputs[0].detach().cpu()
+    logit_labels = torch.argmax(logit, dim=1)
+    #print(str(logit_labels[0]), "\t", ",".join([str(item) for item in all_input_ids[0].cpu()[:10].tolist()]))
+    
+    # What we need to return...
+    # - predicted label, raw logit probability, model L-score, L-score perturbation, Model L-score, cosine distance to nearest proxy in model, euclidean distance to nearest proxy in model
+    # So, we need the models themselves to return some of this stuff...
+    # Models have already been trained. We basically want the model to be equipped with some plugins
+    return logit_labels, logit, None
+
+  def output_setup(self, **kwargs):
+    file_path = kwargs.get("file_name", "deploy-out.txt")
+    self.fobj = open(file_path, "w")
+
+  def output_step(self, logits, features, secondary):
+    self.fobj.write("\n".join([str(item) for item in logits.tolist()])+"\n")
\ No newline at end of file
diff --git a/src/ednaml/core/decorators/__init__.py b/src/ednaml/core/decorators/__init__.py
new file mode 100644
index 00000000..29fa92f1
--- /dev/null
+++ b/src/ednaml/core/decorators/__init__.py
@@ -0,0 +1,56 @@
+import inspect, os
+
+import ednaml
+
+
+REGISTERED_EDNA_COMPONENTS = {}
+
+
+def register(func, functype):
+    try:
+        fname = inspect.getfile(func)
+    except TypeError:
+        fname = os.path.abspath(func.__module__)
+    print("Registering a %s: %s, from file: %s"%(functype, str(func), fname))
+    if fname not in REGISTERED_EDNA_COMPONENTS:
+        REGISTERED_EDNA_COMPONENTS[fname] = {}
+    # TODO also add in some type of lookup for the file basename...?
+    if functype is "model_plugin":
+        if "model_plugin" not in REGISTERED_EDNA_COMPONENTS[fname]:
+            REGISTERED_EDNA_COMPONENTS[fname]["model_plugin"] = []
+        REGISTERED_EDNA_COMPONENTS[fname]["model_plugin"].append(func)
+    else:
+        REGISTERED_EDNA_COMPONENTS[fname][functype] = func
+
+
+def register_crawler(func):
+    register(func, "crawler")
+    return func
+
+def register_model(func):
+    register(func, "model")
+    return func
+
+def register_generator(func):
+    register(func, "generator")
+    return func
+
+def register_storage(func):
+    register(func, "storage")
+    return func
+
+def register_trainer(func):
+    register(func, "trainer")
+    return func
+
+def register_deployment(func):
+    register(func, "deployment")
+    return func
+
+
+def register_model_plugin(func):
+    register(func, "model_plugin")
+    return func
+
+def register_model_pre_plugin(func):
+    pass
diff --git a/src/ednaml/crawlers/__init__.py b/src/ednaml/crawlers/__init__.py
index 1bc90bee..ad20b359 100644
--- a/src/ednaml/crawlers/__init__.py
+++ b/src/ednaml/crawlers/__init__.py
@@ -7,7 +7,7 @@ class Crawler:
 
     classes: Dict[str, int]
 
-    def __init__(self):
+    def __init__(self, logger):
         self.classes = {}
 
     def __verify(self, folder):
@@ -19,20 +19,4 @@ def __verify(self, folder):
             self.logger.info("Found {data_folder}".format(data_folder=folder))
 
 
-# Vehicle Re-ID Crawlers
-from ednaml.crawlers.VeRiDataCrawler import VeRiDataCrawler
-from ednaml.crawlers.VRICDataCrawler import VRICDataCrawler
-from ednaml.crawlers.VehicleIDDataCrawler import VehicleIDDataCrawler
 
-# Carzam crawlers
-from ednaml.crawlers.Cars196DataCrawler import Cars196DataCrawler
-
-# CoLabelCrawlers
-from ednaml.crawlers.VehicleColorCrawler import VehicleColorCrawler
-from ednaml.crawlers.CoLabelIntegratedDatasetCrawler import (
-    CoLabelIntegratedDatasetCrawler,
-)
-
-KnowledgeIntegratedDatasetCrawler = CoLabelIntegratedDatasetCrawler
-
-from ednaml.crawlers.CoLabelCompCarsCrawler import CoLabelCompCarsCrawler
diff --git a/src/ednaml/datareaders/HFReader.py b/src/ednaml/datareaders/HFReader.py
new file mode 100644
index 00000000..e8bd4cf6
--- /dev/null
+++ b/src/ednaml/datareaders/HFReader.py
@@ -0,0 +1,10 @@
+from ednaml.datareaders import DataReader
+from ednaml.generators.HFGenerator import HFGenerator, HFDataset
+from torch.utils.data import Dataset
+from torch.utils.data import DataLoader
+
+
+class HFReader(DataReader):
+    name: str = "HFReader"
+    DATASET: Dataset = HFDataset
+    GENERATOR = HFGenerator
diff --git a/src/ednaml/datareaders/__init__.py b/src/ednaml/datareaders/__init__.py
index 7344cda8..b80d7641 100644
--- a/src/ednaml/datareaders/__init__.py
+++ b/src/ednaml/datareaders/__init__.py
@@ -16,9 +16,6 @@ def __init__(self):
         pass
 
 
-from ednaml.datareaders.VehicleColor import VehicleColor
-from ednaml.datareaders.CompCars import CompCars
-from ednaml.datareaders.VehicleID import VehicleID
-from ednaml.datareaders.VeRi import VeRi
 from ednaml.datareaders.TorchvisionDatareader import TorchvisionDatareader
 from ednaml.datareaders.AlbertReader import AlbertReader
+from ednaml.datareaders.HFReader import HFReader
diff --git a/src/ednaml/deploy/BaseDeploy.py b/src/ednaml/deploy/BaseDeploy.py
index 25395cae..73a64b71 100644
--- a/src/ednaml/deploy/BaseDeploy.py
+++ b/src/ednaml/deploy/BaseDeploy.py
@@ -1,5 +1,5 @@
 import json, logging, os, shutil
-from typing import Dict, List
+from typing import Dict, List, Tuple
 import tqdm
 import torch
 from torch.utils.data import DataLoader
@@ -24,6 +24,7 @@ def __init__(
         self,
         model: ModelAbstract,
         data_loader: DataLoader,
+        epochs: int,
         logger: logging.Logger,
         crawler: Crawler,
         config: EdnaMLConfig,
@@ -36,7 +37,8 @@ def __init__(
        
         self.data_loader = data_loader
         self.logger = logger
-
+        self.epochs = epochs
+        self.global_epoch = 0
         self.global_batch = 0  # Current batch number in the epoch
 
         self.metadata = {}
@@ -50,6 +52,8 @@ def __init__(
             config=json.loads(config.export("json")),
         )
 
+        self.model_is_built = False
+
     def buildMetadata(self, **kwargs):
         for keys in kwargs:
             self.metadata[keys] = kwargs.get(keys)
@@ -59,6 +63,7 @@ def apply(
         step_verbose: int = 5,
         save_directory: str = "./checkpoint/",
         save_backup: bool = False,
+        save_frequency: int = 1,
         backup_directory: str = None,
         gpus: int = 1,
         fp16: bool = False,
@@ -67,6 +72,7 @@ def apply(
     ):
         self.step_verbose = step_verbose
         self.save_directory = save_directory
+        self.save_frequency = save_frequency
         self.backup_directory = None
         self.model_save_name = model_save_name
         self.logger_file = logger_file
@@ -80,9 +86,13 @@ def apply(
         self.gpus = gpus
 
         if self.gpus != 1:
-            raise NotImplementedError()
+            self.logger.warning("Multi-gpu or non-gpu not yet fully supported.")
 
-        self.model.cuda() # moves the model into GPU
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        if self.gpus:
+            #self.model.cuda() # moves the model into GPU
+            self.logger.info("%i GPUs available"%self.gpus)
+        self.model.to(self.device)
 
         self.fp16 = fp16
         # if self.fp16 and self.apex is not None:
@@ -91,48 +101,81 @@ def apply(
     def saveMetadata(self):
         print("NOT saving metadata. saveMetadata() function not set up.")
 
-    def deploy(self, continue_epoch=0, inference = False):
-        self.logger.info("Starting deployment")
-        self.logger.info("Logging to:\t%s" % self.logger_file)
-        if self.config.SAVE.LOG_BACKUP:
-            self.logger.info(
-                "Logs will be backed up to drive directory:\t%s"
-                % self.backup_directory
-            )
-        
-        self.logger.info("Loading model from saved epoch %i" % continue_epoch)
-        if continue_epoch > 0:
-            load_epoch = continue_epoch - 1
-            self.load(load_epoch)
+    def deploy(self, continue_epoch=0, inference = False, ignore_plugins: List[str] = [], execute: bool = True, model_build: bool = None):
+        if model_build or (model_build is None and not self.model_is_built):
+            self.logger.info("Starting deployment")
+            self.logger.info("Logging to:\t%s" % self.logger_file)
+            if self.config.SAVE.LOG_BACKUP:
+                self.logger.info(
+                    "Logs will be backed up to drive directory:\t%s"
+                    % self.backup_directory
+                )
+            
+            #self.logger.info("Loading model from saved epoch %i" % (continue_epoch - 1))
+            if continue_epoch > 0:
+                load_epoch = continue_epoch - 1
+                self.load(load_epoch, ignore_plugins=ignore_plugins)
 
-        if inference:
-            self.model.inference()
+            if inference:
+                self.model.inference()
+            else:
+                self.model.eval()
+            self.model_is_built = True
         else:
-            self.model.eval()
-        self.data_step()
+            if model_build is not None and not model_build:
+                self.logger.info("Skipping model building due to `model_build=False`")
+            elif self.model_is_built:
+                self.logger.info("Skipping model building because model is already built. To force, set the `model_build` flag to True in `ed.deploy`")
+            else:
+                self.logger.info("Skipping model building")
 
-        self.logger.info("Completed deployment task.")
+        if execute:
+            self.logger.info("Setting up plugin hooks. Plugins will fire during:  %s"%self.config.DEPLOYMENT.PLUGIN.HOOKS)
+            self.model.set_plugin_hooks(self.config.DEPLOYMENT.PLUGIN.HOOKS)
+
+            self.logger.info("Executing deployment for  %i epochs" % self.epochs)
+            for epoch in range(self.epochs):
+                self.logger.info("Starting epoch %i"%self.global_epoch)
+                self.model.pre_epoch_hook(epoch=epoch)
+                self.data_step()
+                self.model.post_epoch_hook(epoch=epoch)
+                self.global_epoch = epoch + 1
+                if self.global_epoch % self.save_frequency == 0:
+                    self.save()
+                self.logger.info("Executing end of epoch steps")
+                self.end_of_epoch(epoch=epoch)
+            self.end_of_deployment()
+            self.logger.info("Completed deployment task.")
+        else:
+            self.logger.info("Skipping execution of deployment task due to `execute = False`.")
 
     def data_step(self):   
         with torch.no_grad():
             for batch in tqdm.tqdm(
                 self.data_loader, total=len(self.data_loader), leave=False
             ):    
+                batch = self.move_to_device(batch)
                 feature_logits, features, secondary_outputs = self.deploy_step(batch)
 
                 self.output_step(feature_logits, features, secondary_outputs)
                 # Log Metrics here and inside the model TODO
                 self.global_batch += 1
 
-
+    def move_to_device(self, batch) -> Tuple[torch.Tensor]:
+        return (item.to(self.device) for item in batch)
+        
     def deploy_step(self, batch):   # USER IMPLEMENTS
         batch = tuple(item.cuda() for item in batch)
-        data, labels = batch
+        data, labels = batch    # TODO move plugins here to allow labels as well!!!!!!!!
         feature_logits, features, secondary_outputs = self.model(data)
 
         return feature_logits, features, secondary_outputs
 
-
+    def end_of_epoch(self, epoch: int):
+        pass
+    
+    def end_of_deployment(self):
+        pass
 
     def output_setup(self, **kwargs): # USER IMPLEMENTS; kwargs from config.DEPLOYMENT.OUTPUT_ARGS
         self.logger.info("Warning: No output setup is performed")
@@ -141,7 +184,7 @@ def output_step(self, logits, features, secondary): # USER IMPLEMENTS, ALSO, NEE
         if self.global_batch % self.config.LOGGING.STEP_VERBOSE == 0:
             self.logger.info("Warning: No output is generated at step %i"%self.global_batch)
 
-    def load(self, load_epoch):
+    def load(self, load_epoch, ignore_plugins: List[str] = []):
         self.logger.info(
             "Loading a model from saved epoch %i"
             % (load_epoch)
@@ -162,4 +205,65 @@ def load(self, load_epoch):
         self.model.load_state_dict(torch.load(model_load_path))
         self.logger.info(
             "Finished loading model state_dict from %s" % model_load_path
-        )
\ No newline at end of file
+        )
+        # Here, we will need to get a list of pickled or serialized plugin objects, then load them into a dictionary, then pass them into model
+        # YES!!!!
+        plugin_load = self.model_save_name + "_plugins.pth"
+
+        if self.save_backup:
+            self.logger.info(
+                "Looking for model plugins from drive backup."
+            )
+            plugin_load_path = os.path.join(self.backup_directory, plugin_load)
+        else:
+            self.logger.info(
+                "Looking for model plugins from local backup."
+            )
+            plugin_load_path = os.path.join(self.save_directory, plugin_load)
+
+        if os.path.exists(plugin_load_path):
+            
+            self.model.loadPlugins(plugin_load_path, ignore_plugins=ignore_plugins)
+            self.logger.info(
+                "Loaded plugins from %s"%plugin_load_path
+            )
+        else:
+            self.logger.info(
+                "No plugins found at %s"%plugin_load_path
+            )
+
+
+    def save(self, save_epoch: int = None):
+        if save_epoch is None:
+            save_epoch = self.global_epoch
+        self.logger.info("Performing save at epoch %i"%save_epoch)
+        self.logger.info("Saving model plugins.")
+        PLUGIN_SAVE_NAME = self.model_save_name + "_plugins.pth"
+        # here, we save only the plugins...
+        # Do not need model save, or training save, or log save...
+        
+        # A ModelAbstract can save its own plugins...?
+        # No, we get back a plugin serialization object that we use to , well, save the plugin
+        # Later, during plugin instantiation, we will need to write the code to load the plugin's save code...
+        plugin_save = self.model.savePlugins()
+
+        if len(plugin_save) > 0:
+            torch.save(plugin_save, os.path.join(
+                self.save_directory, PLUGIN_SAVE_NAME
+            ))
+            self.logger.info("Saved plugins: %s"%str(plugin_save.keys()))
+
+            if self.save_backup:
+                shutil.copy2(
+                    os.path.join(self.save_directory, PLUGIN_SAVE_NAME),
+                    self.backup_directory,
+                )
+
+                self.logger.info(
+                    "Performing drive backup of model plugins"
+                )
+        else:
+            self.logger.info("No plugins to save")
+
+       
+        
\ No newline at end of file
diff --git a/src/ednaml/deploy/HFDeploy.py b/src/ednaml/deploy/HFDeploy.py
new file mode 100644
index 00000000..8476c74e
--- /dev/null
+++ b/src/ednaml/deploy/HFDeploy.py
@@ -0,0 +1,32 @@
+
+from ednaml.deploy.BaseDeploy import BaseDeploy
+
+class HFDeploy(BaseDeploy):
+  def deploy_step(self, batch):
+    #batch = tuple(item.cuda() for item in batch)
+    (
+        all_input_ids,
+        all_attention_mask,
+        all_token_type_ids,
+        all_masklm,
+        all_annotations,
+        all_labels
+    ) = batch
+    outputs = self.model(
+        all_input_ids,
+        token_type_ids=all_token_type_ids,
+        attention_mask=all_attention_mask,
+        output_attentions = True,
+        output_hidden_states = True,
+        secondary_inputs=all_annotations       # NOT for HFTrainer! because it only expects specific inputs!
+    )
+
+    logits = outputs[0]
+    features = outputs[1]
+    secondaries = outputs[2]
+    
+    return logits, features, secondaries
+  def output_setup(self, **kwargs):
+    pass
+  def output_step(self, logits, features, secondary):
+    pass
\ No newline at end of file
diff --git a/src/ednaml/deploy/__init__.py b/src/ednaml/deploy/__init__.py
index 0cce3061..a71b4555 100644
--- a/src/ednaml/deploy/__init__.py
+++ b/src/ednaml/deploy/__init__.py
@@ -1,1 +1,2 @@
-from ednaml.deploy.BaseDeploy import BaseDeploy
\ No newline at end of file
+from ednaml.deploy.BaseDeploy import BaseDeploy
+from ednaml.deploy.HFDeploy import HFDeploy
\ No newline at end of file
diff --git a/src/ednaml/generators/AlbertGenerator.py b/src/ednaml/generators/AlbertGenerator.py
index fa97fc86..8107385b 100644
--- a/src/ednaml/generators/AlbertGenerator.py
+++ b/src/ednaml/generators/AlbertGenerator.py
@@ -30,7 +30,10 @@ def __init__(self, logger, dataset, mode, transform=None, **kwargs):
                 shutil.rmtree(self.shardpath)
             else:
                 self.shards_exist = True
-        os.makedirs(self.shardpath, exist_ok=True)
+                self.logger.debug("Shards already exist and `shard_replace` is False")
+        if self.shardcache:
+            self.logger.debug("Creating shardpath %s"%self.shardpath)
+            os.makedirs(self.shardpath, exist_ok=True)
         
 
 
@@ -118,8 +121,8 @@ def shardget(self, idx):
             self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
             if self.data_shuffle:
                 random.shuffle(self.shard_internal_shuffle)
-        shardindex = idx % self.current_shardsize
-        return self.sharded_dataset[self.shard_internal_shuffle[shardindex]]
+        #shardindex = idx % self.current_shardsize
+        return self.sharded_dataset[self.shard_internal_shuffle[self.getcount]]
 
     def sharded_convert_to_features(self, dataset, tokenizer, maxlen):
         features = []
diff --git a/src/ednaml/generators/ClassificationGenerator.py b/src/ednaml/generators/ClassificationGenerator.py
index be608cfc..03818e39 100644
--- a/src/ednaml/generators/ClassificationGenerator.py
+++ b/src/ednaml/generators/ClassificationGenerator.py
@@ -39,36 +39,21 @@ def load(self, img):
 
 
 class ClassificationGenerator(ImageGenerator):
-    def __init__(self, logger = None, gpus=1, **kwargs):
-        """Data generator for training and testing. Works with the <>. Should work with any crawler working on VeRi-like data. Not yet tested with VehicleID. Only  use with VeRi.
 
-        Generates batches of batch size CONFIG.TRANSFORMATION.BATCH_SIZE, with CONFIG.TRANSFORMATION.INSTANCE unique ids. So if BATCH_SIZE=36 and INSTANCE=6, then generate batch of 36 images, with 6 identities, 6 image per identity. See arguments of setup function for INSTANCE.
-
-        Args:
-            gpus (int): Number of GPUs
-            i_shape (int, int): 2D Image shape
-            normalization_mean (float): Value to pass as mean normalization parameter to pytorch Normalization
-            normalization_std (float): Value to pass as std normalization parameter to pytorch Normalization
-            normalization_scale (float): Value to pass as scale normalization parameter. Not used.
-            h_flip (float): Probability of horizontal flip for image
-            t_crop (bool): Whether to include random cropping
-            rea (bool): Whether to include random erasing augmentation (at 0.5 prob)
-
-        """
-        super().__init__(logger, gpus=gpus, **kwargs)
 
     def buildDataset(
         self, datacrawler, mode: str, transform: List[object], **kwargs
     ) -> TorchDataset:
-        if mode == "train":
+        if mode in ["train", "val", "test"]:
             return ClassificationDataset(
-                datacrawler.metadata[mode]["crawl"], transform, **kwargs
-            )
-        elif mode == "test":
+                datacrawler.metadata[mode]["crawl"],
+                transform,
+                **kwargs
+            )    
+        elif mode == "testval":
             # For testing, we combine images in the query and testing set to generate batches
             return ClassificationDataset(
-                datacrawler.metadata["val"]["crawl"]
-                + datacrawler.metadata[mode]["crawl"],
+                datacrawler.metadata["val"]["crawl"]+datacrawler.metadata["test"]["crawl"],
                 transform,
                 **kwargs
             )
@@ -92,11 +77,11 @@ def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
                 num_workers=self.workers,
                 collate_fn=self.collate_simple,
             )
-        elif mode == "test":
+        elif mode in ["test", "val"]:
             return TorchDataLoader(
                 dataset,
                 batch_size=batch_size * self.gpus,
-                shuffle=kwargs.get("shuffle", True),
+                shuffle=kwargs.get("shuffle", False),
                 num_workers=self.workers,
                 collate_fn=self.collate_simple,
             )
@@ -112,7 +97,7 @@ def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
             raise NotImplementedError()
 
     def getNumEntities(self, datacrawler, mode, **kwargs):
-        if mode in ["train", "test", "full"]:
+        if mode in ["train", "test", "full", "val"]:
             label_dict = {
                 item: {"classes": datacrawler.metadata[mode]["classes"][item]}
                 for item in [kwargs.get("classificationclass", "color")]
diff --git a/src/ednaml/generators/HFGenerator.py b/src/ednaml/generators/HFGenerator.py
new file mode 100644
index 00000000..adef3590
--- /dev/null
+++ b/src/ednaml/generators/HFGenerator.py
@@ -0,0 +1,615 @@
+from glob import glob
+import random
+import torch, os, shutil
+from torch.utils.data import TensorDataset, IterableDataset, Dataset
+import numpy as np, h5py
+from tqdm import tqdm
+class HFDataset(Dataset):
+    def __init__(self, logger, dataset, mode, transform=None, **kwargs):
+        """Initializes the HFDataset object.
+
+       Args:
+            logger (_type_): The logger object
+            dataset (_type_): The crawler's metadata object
+            mode (_type_): Mode to extract crawler contents
+            transform (_type_, optional): Unused. Defaults to None.
+
+        Kwargs (DO NOT CHANGE IN CONFIG):
+            tokenizer: The tokenizer object
+
+        COnfig Args (provided through DATASET_KWARGS):
+            data_shuffle: Whether to shuffle the samples.
+
+            cache: Cache preprocessed samples in an h5 file. Samples are accessed ad-hoc.
+            memcache: Cache preprocessed samples in-memory. Samples are accessed ad-hoc.
+            shardcache: Cache preprocessed samples in shards on disk. Samples are limited to
+                shard-based access.
+
+            cache_replace: Whether to replace the h5 cache, if it exists.
+            cachepath: Constructor for the directory where the cache is kept. This
+                is appended with the `mode`. So, for training, a `cachepath` of 
+                "datashard-artifacts" becomes "datashard-artifacts-train"
+            cachename: The name for the cache file stored inside the `cachepath`.
+                So a cachepath of name "hf-cache" means the following file is created:
+                    "datashard-artifacts-train/hf-cache.h5" 
+
+
+            shardsize: Number of samples in each shard. This in turn determines 
+                how much RAM is dedicated to hosting a shard in-memory when it is loaded.
+            shard_replace: If shards already exist, whether to replace them.
+            shardpath: Constructor for the directory where shards are kept. This
+                is appended with the `mode`. So, for training, a `shardpath` of 
+                "datashard-artifacts" becomes "datashard-artifacts-train"
+            shard-name: The name for each shard stored inside the `shardpath`.
+                This value is appended with the shard number. So a training shard
+                with index 5 and name "hf-shard" is stored at:
+                    "datashard-artifacts-train/hf-shard5.pt" 
+
+            maxlen: The maximum length of the text to encode. Usually 512.
+            mlm_probability: The probability of masking if using word and token masking.
+                Usually around 0.15
+            masking (bool): Whether to use any masking. Specific masking options are next. Can also be `train_only` for masking during training only.
+            token_mask: Whether to use random token masking, with `mlm_probability`
+            word_mask: Whether to use random word masking, with `mlm_probability`
+            keytoken_mask: Whether to use specific token masking, with `keytokens`
+            keyword_mask: Whether to use specific word masking, with `keywords`
+            keywords: What words to mask
+            keytokens: What tokens to mask
+            stopword_mask: Whether to use a stopword mask.
+            label_idxs: List of indices for labels. Defaults to [1].
+            annotation_idxs: List of indices for inputs. Defaults to [].
+
+
+            
+        """
+        self.dataset = dataset  # list of tuples or dict (text, *text, *labels)
+        self.logger = logger
+        self.file_len = len(dataset)
+        self.data_shuffle = kwargs.get("data_shuffle", True)
+
+        # Cacheing options
+        self.cache = kwargs.get("cache", False)
+        self.memcache = kwargs.get("memcache", False)
+        self.shardcache = kwargs.get("shardcache", False)
+
+        if int(self.cache) + int(self.memcache) + int(self.shardcache) > 1:
+            raise RuntimeError("Use only one cache type")
+
+        self.word_masking = kwargs.get("word_mask", False)
+        self.keyword_masking = kwargs.get("keyword_mask", True)
+        self.token_masking = kwargs.get("token_mask", False)
+        self.keytoken_masking = kwargs.get("keytoken_mask", False)
+        self.stopword_masking = kwargs.get("stopword_mask", False)
+
+        self.label_idxs = kwargs.get("label_idxs", [1])
+        self.annotation_idxs = kwargs.get("annotation_idxs", self.label_idxs)
+        
+        
+        # Tokenizer and lengths
+        self.tokenizer = kwargs.get("tokenizer")
+        self.maxlen = kwargs.get("maxlen")
+        self.mlm = kwargs.get("mlm_probability", 0.2)
+        self.masking = kwargs.get("masking", mode=="train")
+        if self.masking == "train_only":
+            self.masking = mode=="train"
+
+
+        self.keywords = []
+        if self.keyword_masking:
+            self.keywords = kwargs.get("keywords", [])
+        self.keytokens = []
+        if self.keytoken_masking:
+            self.keytokens = kwargs.get("keytokens", [])
+        # Store indices corresponding to keytokens
+        vocab = self.tokenizer.get_vocab()
+        self.keytokens = {vocab[item]:1 for item in self.keytokens if item in vocab}
+        self.keytokens_tensor = torch.tensor(list(self.keytokens.keys())).unsqueeze(0).T
+
+
+        # Cache options
+        self.cache_replace = kwargs.get("cache_replace", False)
+        self.cachepath = kwargs.get("cachepath", "artifacts")
+        self.cachepath = mode+"-"+self.cachepath
+        self.cachename = kwargs.get("cachename", "h5-cache") + "-"  #the dash
+        if self.cache:
+            self.logger.debug("[Mode `{mode}`] Will look in path [{path}] for shards `{shards}[e].pt`".format(path=self.cachepath, shards=self.cachename, mode=mode))
+            self.base_cachepath = os.path.join(self.cachepath, self.cachename) + ".h5"
+            self.cache_exist = False
+            if os.path.exists(self.base_cachepath):
+                if self.cache_replace:
+                    self.logger.debug("Deleting existing cache")
+                    os.remove(self.cachepath)
+                else:
+                    self.cache_exist = True
+                    self.logger.debug("Cache already exists and `cache_replace` is False")
+            else:
+                self.logger.debug("Cache does not exist and will be created.")
+                self.logger.debug("Creating cachepath %s"%self.cachepath)
+                os.makedirs(self.cachepath, exist_ok=True)
+                with h5py.File(self.base_cachepath, "w") as hfile:
+                    hfile.create_dataset(name="all_input_ids", shape=(len(self.dataset),self.maxlen))
+            
+
+
+        # Shardcache options
+        self.shardsize = kwargs.get("shardsize", 0)
+        self.shardsaveindex = 0
+        self.shard_replace = kwargs.get("shard_replace", False)
+        self.shardpath = kwargs.get("shardpath", "datashard-artifacts")
+        self.shardpath = mode+"-"+self.shardpath
+        self.shardname = kwargs.get("shardname", "fnc-filtermask-shard") + "-"  #the dash
+        if self.shardcache:
+            self.logger.debug("[Mode `{mode}`] Will look in path [{path}] for shards `{shards}[e].pt`".format(path=self.shardpath, shards=self.shardname, mode=mode))
+        self.base_shardpath = os.path.join(self.shardpath, self.shardname)
+        self.shards_exist = False
+        if os.path.exists(self.base_shardpath + "0.pt"):
+            if self.shard_replace:
+                self.logger.debug("Deleting existing shards")
+                shutil.rmtree(self.shardpath)
+            else:
+                self.shards_exist = True
+                self.logger.debug("Shards already exist and `shard_replace` is False")
+        else:
+          self.logger.debug("Shards do not exist and will be created.")
+        if self.shardcache:
+            self.logger.debug("Creating shardpath %s"%self.shardpath)
+            os.makedirs(self.shardpath, exist_ok=True)
+        
+        
+        
+        
+
+        # This is for memcache and shardcache
+        self.getcount = 0
+        self.shardcount = 0
+        if self.memcache or self.cache:
+            # Checks when we have processed entire dataset, refreshes mask ids when needed
+            self.refresh_flag = len(dataset)
+        if self.shardcache:
+            # We need to refresh mask ids each time we load a shard.
+            # This is ONLY for randomized masking, i.e. with MLM probability.
+            self.refresh_flag = self.shardsize
+            # this shuffles the internal shard index, so that we get a random example, instead of in order
+            self.shard_internal_shuffle = []
+            # This shuffles the shard index, so that each epoch, we load shards in different orders. We will set this up after setting up shards...
+            self.shard_shuffle = []
+        
+        self.getter = self.uncachedget  # for each entry, directly convert to features and return
+        if self.cache:
+            self.getter = self.diskget
+        elif self.memcache:
+            self.getter = self.memget
+        elif self.shardcache:
+            self.getter = self.shardget
+
+        if self.cache or self.memcache:
+        # The actual cache-ing
+            self.logger.debug("Started mem caching")
+            self.input_length_cache = []
+            self.logger.debug("Converting to features")
+            # For memcache, save dataset inside self.cached_dataset.
+            # Same for cache. Difference being, cached_dataset is either in memory or on disk!
+            self.convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen, cache=True)
+            self.logger.debug("Masking")
+            self.memcached_dataset = self.refresh_mask_ids(cache=True)
+
+        if self.shardcache:
+            self.input_length_cache = []
+            if not self.shards_exist:
+                self.logger.debug("Generating shards")
+                self.shardsaveindex = self.sharded_convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen)    # save shards and get numshards
+            else:
+                self.shardsaveindex = len(glob(os.path.join(self.shardpath, "*.pt")))   # TODO Bug fix if files do not have consistent numbering
+            if self.shardsaveindex < 1:
+                raise ValueError("`shardsaveindex` is {val}, which is less than permissible minimum value of 1".format(val=self.shardsaveindex))
+            self.logger.debug("Obtained %i shards"%self.shardsaveindex)
+            self.shard_load_index = 0   # self.shardsaveindex is the maximum number of shards
+            self.shard_shuffle = list(range(self.shardsaveindex))    # count started from 0
+            if self.data_shuffle:
+                self.logger.debug("Shuffling shard load order")
+                random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+            self.current_shardsize = len(self.sharded_dataset)
+            if self.masking:
+                self.logger.debug("Refreshing token masks for loaded shard")    # TODO
+                self.sharded_dataset = self.sharded_refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+
+
+
+    def shardget(self, idx):
+        """Given an index, retrieve a single sample from the post-processed
+        dataset.
+
+        We retrieve from shards, so we need to track the current shard we are
+        working with, as well as which sample in the shard we are retrieving.
+
+        THus, given an index, we do not use that index directly. Instead,
+        HFGenerator tracks samples itself by using `self.shard_shuffle` 
+        and `self.internal_shard_shuffle`. The former tracks which shards
+        have been accessed already, and the latter tracks which samples
+        in the current shard have been tracked already.
+
+        `self.getcount` tracks the current index.
+
+        Shuffling if controlled with the `data_shuffle` argument during
+        initialization of the HFDataset.
+
+        Args:
+            idx (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        # get entry from already loaded shardcache. so idx is incremented one by one -- No shuffling in data loader.
+        # This is a design choice, and we can't really do anything if user does shuffle + sharding
+        self.last_idx = idx
+        response = self.sharded_dataset[self.shard_internal_shuffle[self.getcount]]
+        self.getcount += 1
+
+        if self.getcount == self.current_shardsize:   # we have exhausted examples in this shard
+            self.getcount = 0
+            self.shard_load_index += 1                  # increment the shard index that we will load
+            if self.shard_load_index == self.shardsaveindex: # we have processed all shards.
+                self.shard_load_index = 0
+                if self.data_shuffle:
+                    random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+
+            self.current_shardsize = len(self.sharded_dataset)
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+            if self.masking:
+                self.sharded_dataset = self.sharded_refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+
+        return response
+
+    def sharded_convert_to_features(self, dataset, tokenizer, maxlen):
+        features = []
+        shardsaveindex = 0
+        input_length_cache = []
+        input_word_length_cache = []
+        self.logger.info("Generating shards")
+        pbar = tqdm(total=int(self.file_len/self.shardsize)+1)
+        # TODO
+        # We basically need to deal with having multiple inputs and multiple labels...
+        # Ok, so, we already know the structure of the tuple that dataset gets
+        # Then, all we need is in config, specify additional_input_idx --> list | None
+        # if None, we don't need to do anything
+        # if a list, then our input is actually a tuple
+        # So, if a list, what do we do...?
+        # 
+        #
+        #
+        #
+        #
+        for idx, sample in enumerate(dataset):  # This is a tuple with (text, ..., stuff)
+            # Identify the indices of specific keywords to mask
+            #if self.keyword_masking or self.word_masking:
+            # NOTE: we assume that text is at index 0!!!!!
+            word_tokens = sample[0].lower().split(" ")
+            encoded_word_length = len(word_tokens)
+            
+            if self.masking and self.keyword_masking:
+                keyword_idx = []
+                for widx,item in enumerate(word_tokens):
+                    if sum([1 if kword in item else 0 for kword in self.keywords]) > 0:
+                        keyword_idx.append(widx) # i.e. find word index for each keyword, if it exists
+
+            encoded = self.tokenizer(sample[0], return_tensors="pt", padding="max_length", max_length = maxlen, truncation = True, return_length = True)
+            encoded_token_length = torch.sum(encoded["attention_mask"])
+            encoded_word_ids = encoded.word_ids(0)
+            input_length_cache.append(encoded_token_length)
+            input_word_length_cache.append(encoded_word_length)
+            # create a list of lists. Each sublist is a mask for each keyword. So, we can AND all sublists to get the overall mask.
+            # Edge cases -- no keywords: 
+            merged_mask = encoded["attention_mask"]
+            if self.masking and self.keyword_masking:
+                match_idxs = torch.LongTensor([[word_idx != idx_of_keyword for word_idx in encoded_word_ids] for idx_of_keyword in keyword_idx])
+                if match_idxs.shape[0] > 0: #i.e. we have a mask for keywords, so we will and everything
+                    match_idxs = torch.all(match_idxs, dim=0, keepdim=True)
+                    merged_mask = torch.where(match_idxs == 0, match_idxs, merged_mask)
+
+            if self.masking and self.keytoken_masking:
+                # Dict with relevant keyword idxs...self.keytokens 
+                # First set 1 for all tokens, and 0 for masking tokens
+                match_idxs = encoded["input_ids"].repeat((len(self.keytokens),1))
+                match_idxs = torch.all(match_idxs != self.keytokens_tensor, dim=0, keepdim=True)
+                merged_mask = torch.where(match_idxs == 0, match_idxs, merged_mask)
+
+            
+            features.append(
+                (encoded["input_ids"], merged_mask, encoded["token_type_ids"], 
+                    encoded_token_length, encoded_word_length, encoded_word_ids, 
+                    [sample[label_idx] for label_idx in self.label_idxs],   # list of labels
+                    [sample[annot_idx] for annot_idx in self.annotation_idxs])   # list of secondary inputs. Defaults to empty.
+            )
+
+            if (idx+1)%self.shardsize == 0: # we need to save this shard.
+                # TODO need to handle multiple labels, later...
+                all_input_ids = torch.cat([f[0] for f in features])
+                # TODO
+                # Secondary inputs are other numeric inputs
+                # We check if secondary is provided by checking what is in  in features[7] (last element)
+                # if None, then secondary inputs are not provided. We can make it a matrix of zeros
+                # # if true, the secondary inputs are provided, in a list of numbers
+                # Then we convert then to a nxd matrix (n = shardsize, d = number of secondary inputs)
+                #all_secondary_inputs = None
+                all_attention_mask = torch.cat([f[1] for f in features])
+                all_token_type_ids = torch.cat([f[2] for f in features])
+                all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+                all_word_lens = torch.tensor([f[4] for f in features], dtype=torch.long)
+                all_word_ids = torch.tensor(np.nan_to_num(np.array([f[5] for f in features], dtype=float),nan=-1),dtype=torch.long)
+                all_labels = torch.tensor([f[6] for f in features], dtype=torch.long)   # TODO -- implement multiple labels as a list, then squash here to one dim if single label.
+                all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+                if len(features[0][7]) > 0:
+                    all_annotations = torch.tensor([f[7] for f in features])
+                else:
+                    all_annotations = torch.zeros((self.shardsize,1))
+
+                # Propagate input masking to output masking
+                for midx in range(all_attention_mask.shape[0]):
+                    masked_index = torch.where(all_attention_mask[midx][:input_length_cache[midx]]==0)[0]
+                    all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+                input_length_cache = []
+                input_word_length_cache = []
+
+                # SAVE HERE
+                self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations),
+                        shard_index = shardsaveindex)
+                shardsaveindex += 1
+                features = []
+                pbar.update(1)
+        # final shard...
+        if len(features) > 0:
+            all_input_ids = torch.cat([f[0] for f in features])
+            all_attention_mask = torch.cat([f[1] for f in features])
+            all_token_type_ids = torch.cat([f[2] for f in features])
+            all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+            all_word_lens = torch.tensor([f[4] for f in features], dtype=torch.long)
+            all_word_ids = torch.tensor(np.nan_to_num(np.array([f[5] for f in features], dtype=float),nan=-1),dtype=torch.long)
+            all_labels = torch.tensor([f[6] for f in features], dtype=torch.long)
+            all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+            if len(features[0][7]) > 0:
+                    all_annotations = torch.tensor([f[7] for f in features])
+            else:
+                all_annotations = torch.zeros((len(features),1))
+
+            # Propagate input masking to output masking
+            for midx in range(all_attention_mask.shape[0]):
+                masked_index = torch.where(all_attention_mask[midx][:input_length_cache[midx]]==0)[0]
+                all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+            input_length_cache = []
+
+            # SAVE HERE
+            self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations),
+                    shard_index = shardsaveindex)
+            shardsaveindex += 1
+            features = []
+            pbar.update(1)
+        pbar.close()
+        return shardsaveindex
+
+    def save_shard(self, shard: TensorDataset, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        torch.save(shard, shard_save_path)
+
+    def load_shard(self, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        self.logger.debug("Loading shard %s"%shard_save_path)
+        return torch.load(shard_save_path)
+
+    def memget(self, idx):
+        self.getcount+=1
+        if self.getcount > self.refresh_flag and self.masking:
+            self.memcached_dataset = self.refresh_mask_ids()
+            self.getcount = 0
+        return self.memcached_dataset[idx]
+
+    def refresh_mask_ids(self):
+        self.logger.debug("Refreshing mask ids")
+        if self.masking:
+            self.mask_ids = self.build_mask_ids(self.input_length_cache)
+            
+            all_attention_mask = self.all_attention_mask.clone()
+            all_masklm = self.all_masklm.clone()
+
+            for idx in range(self.refresh_flag):
+                all_attention_mask[idx][self.mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+                all_masklm[idx][self.mask_ids[idx]] = self.all_input_ids[idx][self.mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+
+            return TensorDataset(self.all_input_ids, all_attention_mask, self.all_token_type_ids, all_masklm, self.all_lens, self.all_labels)
+        else:
+            return TensorDataset(self.all_input_ids, self.all_attention_mask, self.all_token_type_ids, self.all_masklm, self.all_lens, self.all_labels)
+
+    def build_mask_ids(self, input_length_cache):
+        # for each element, we get a set of indices that are randomly selected...
+        # Also, -2 and +1 take care of [cls] and [sep] not being masked
+        return [(torch.randperm(inplength-2)+1)[:int(inplength*self.mlm)]  for inplength in input_length_cache]
+    
+    def build_whole_word_mask_ids(self, input_length_cache):
+        # for each element, we get a set of indices that are randomly selected...
+        # Also, -2 and +1 take care of [cls] and [sep] not being masked
+        return [(torch.randperm(word_length)+1)[:int(word_length*self.mlm)]  for word_length in input_length_cache]
+        
+
+        # So, given word_len, we can extract which ids need to be masked
+        # But, now we need a way to, given the ID that needs to be masked, get back the original words.
+    
+    def sharded_refresh_mask_ids(self, sharded_dataset: TensorDataset):
+        self.logger.debug("Refreshing mask ids")
+        #                   0               1                   2                 3           4           5             6            7............8
+        #TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations)
+                      
+        if self.token_masking or self.word_masking:
+            #merged_masklm =  torch.stack([shard[3] for shard in sharded_dataset])
+            all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+            all_attention_mask = torch.stack([shard[1] for shard in sharded_dataset])
+            if self.token_masking:
+                self.logger.debug("Performing random token masking")
+                mask_ids = self.build_mask_ids(torch.stack([shard[4] for shard in sharded_dataset]))
+                
+                
+                #all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+                all_input_ids = torch.stack([shard[0] for shard in sharded_dataset])
+
+                for idx in range(self.current_shardsize):
+                    all_attention_mask[idx][mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+                    all_masklm[idx][mask_ids[idx]] = all_input_ids[idx][mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+                #merged_masklm *= all_masklm
+            if self.word_masking:
+                self.logger.debug("Performing random word masking")
+                # TODO
+                masking_words = self.build_whole_word_mask_ids(torch.stack([shard[5] for shard in sharded_dataset]))
+                
+                #all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+                all_input_ids = torch.stack([shard[0] for shard in sharded_dataset])
+                all_word_ids = torch.stack([shard[6] for shard in sharded_dataset])
+                for idx in range(self.current_shardsize):
+                
+                    match_idxs = torch.LongTensor([[word_idx != idx_of_keyword for word_idx in all_word_ids[idx]] for idx_of_keyword in masking_words[idx]])
+                    if match_idxs.shape[0] > 0: #i.e. we have a mask for keywords, so we will and everything
+                        match_idxs = torch.all(match_idxs, dim=0, keepdim=True)
+                        merged_mask = torch.where(match_idxs[0]==0)[0]
+                        
+                        all_attention_mask[idx][merged_mask] = 0    # Set the masking words to 0, so we do not attend during prediction
+                        all_masklm[idx][merged_mask] = all_input_ids[idx][merged_mask] # Set the masking labels for these to the actual word index from all_input_ids
+                
+                #merged_masklm *= all_masklm
+            return TensorDataset(all_input_ids, all_attention_mask, torch.stack([shard[2] for shard in sharded_dataset]), all_masklm, 
+                torch.stack([shard[4] for shard in sharded_dataset]), 
+                torch.stack([shard[5] for shard in sharded_dataset]), 
+                torch.stack([shard[6] for shard in sharded_dataset]), 
+                torch.stack([shard[7] for shard in sharded_dataset]),
+                torch.stack([shard[8] for shard in sharded_dataset]))
+        else:
+            return sharded_dataset
+
+
+    def convert_to_features(self, dataset, tokenizer, maxlen, cache=False):
+        """Preprocess the dataset into text features. If cache is True, save the preprocessed data into a local h5 file.
+        Else, save in memory.
+
+        Args:
+            dataset (_type_): _description_
+            tokenizer (_type_): _description_
+            maxlen (_type_): _description_
+            cache (bool, optional): _description_. Defaults to False.
+        """
+        features = []
+        self.input_length_cache = []
+        for idx, sample in enumerate(dataset):
+            tokens = self.tokenizer.tokenize(sample[0])
+            if len(tokens) > maxlen - 2:
+                tokens = tokens[0:(maxlen - 2)]
+            # An easy way potentially -- get the index of each word in the token set...
+            # Then check which word is in keyword
+            # Then mask out that word...????????
+            finaltokens = ["[CLS]"]
+            token_type_ids = [0]
+            for token in tokens:
+                finaltokens.append(token)
+                token_type_ids.append(0)
+            finaltokens.append("[SEP]")
+            token_type_ids.append(0)
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(finaltokens)
+            attention_mask = [1]*len(input_ids)
+            input_len = len(input_ids)
+            self.input_length_cache.append(len(input_ids))
+            while len(input_ids) < maxlen:
+                input_ids.append(0)
+                attention_mask.append(0)
+                token_type_ids.append(0) 
+
+            assert len(input_ids) == maxlen
+            assert len(attention_mask) == maxlen
+            assert len(token_type_ids) == maxlen
+        
+            features.append(
+                (input_ids, attention_mask, token_type_ids, input_len, *sample[1:])
+            )
+
+        self.all_input_ids = torch.tensor([f[0] for f in features], dtype=torch.long)
+        self.all_attention_mask = torch.tensor([f[1] for f in features], dtype=torch.long)
+        self.all_token_type_ids = torch.tensor([f[2] for f in features], dtype=torch.long)
+        self.all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+        self.all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+        self.all_masklm = -1*torch.ones(self.all_input_ids.shape, dtype=torch.long)
+                
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.getter(idx)
+
+    def uncachedget(self, idx):
+        return 
+
+
+
+
+
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.generators import TextGenerator
+class HFGenerator(TextGenerator):
+    """HFGenerator is a generic text generator for HuggingFace-tokenizers
+    and designed to work with HuggingFace style transformers.
+
+    It can use any tokenizer that follows HuggingFace API.
+
+    Usage:
+        TODO
+    """
+
+    def build_transforms(self, transform, mode, **kwargs):  #<-- generator kwargs:
+        """Use the GENERATOR_KWARGS to identify the tokenizer
+        class. For now, only built-in tokenizers are supported.
+
+        Args:
+            transform (_type_): Unused. It is necessary here as a dummy argument in 
+                constructor for the core ednaml.generator API.
+            mode (_type_): Whether to use train or test mode. Useful when selecting
+                the correct subset from crawler.
+        """
+        from ednaml.utils import locate_class
+        print("Building Transforms")
+        tokenizer = kwargs.get("tokenizer", "HFAutoTokenizer")
+        self.tokenizer = locate_class(package="ednaml", subpackage="utils", classpackage=tokenizer, classfile="tokenizers")
+        self.tokenizer = self.tokenizer(**kwargs) # vocab_file, do_lower_case, spm_model_file
+
+    
+    def buildDataset(self, crawler, mode, transform, **kwargs): #<-- dataset args:
+        print("Building Dataset")
+        return HFDataset(self.logger, crawler.metadata[mode]["crawl"], mode, tokenizer = self.tokenizer, **kwargs) # needs maxlen, memcache, mlm_probability
+
+    def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
+        print("Building Dataloader")
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size*self.gpus,
+                                            shuffle=False, num_workers = self.workers, 
+                                        collate_fn=self.collate_fn)
+
+    def getNumEntities(self, crawler, mode, **kwargs):  #<-- dataset args
+        label_dict = {
+            item: {"classes": crawler.metadata[mode]["classes"][item]}
+            for item in kwargs.get("classificationclass", ["color"])
+        }
+        return LabelMetadata(label_dict=label_dict)
+
+    def collate_fn(self, batch):
+        all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations = map(torch.stack, zip(*batch))
+        max_len = max(all_lens).item()
+        all_input_ids = all_input_ids[:, :max_len]
+        all_attention_mask = all_attention_mask[:, :max_len]
+        all_token_type_ids = all_token_type_ids[:, :max_len]
+        all_masklm = all_masklm[:, :max_len]
+        # Right now, we don't care that much about all_word_ids or all_word_lens
+
+        return all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_annotations, all_labels
+
+
+
+
diff --git a/src/ednaml/generators/HFGeneratorOld.py b/src/ednaml/generators/HFGeneratorOld.py
new file mode 100644
index 00000000..64f6753e
--- /dev/null
+++ b/src/ednaml/generators/HFGeneratorOld.py
@@ -0,0 +1,609 @@
+from glob import glob
+import random
+import torch, os, shutil
+from torch.utils.data import TensorDataset, Dataset
+import numpy as np
+from tqdm import tqdm
+class HFDataset(Dataset):
+    def __init__(self, logger, dataset, mode, transform=None, **kwargs):
+        """Initializes the HFDataset object.
+
+       Args:
+            logger (_type_): The logger object
+            dataset (_type_): The crawler's metadata object
+            mode (_type_): Mode to extract crawler contents
+            transform (_type_, optional): Unused. Defaults to None.
+
+        Kwargs (DO NOT CHANGE IN CONFIG):
+            tokenizer: The tokenizer object
+
+        COnfig Args (provided through DATASET_KWARGS):
+            data_shuffle: Whether to shuffle the samples.
+
+            cache: Cache preprocessed samples in an h5 file. Samples are accessed ad-hoc.
+            memcache: Cache preprocessed samples in-memory. Samples are accessed ad-hoc.
+            shardcache: Cache preprocessed samples in shards on disk. Samples are limited to
+                shard-based access.
+
+            cache_replace: Whether to replace the h5 cache, if it exists.
+            cachepath: Constructor for the directory where the cache is kept. This
+                is appended with the `mode`. So, for training, a `cachepath` of 
+                "datashard-artifacts" becomes "datashard-artifacts-train"
+            cachename: The name for the cache file stored inside the `cachepath`.
+                So a cachepath of name "hf-cache" means the following file is created:
+                    "datashard-artifacts-train/hf-cache.h5" 
+
+
+            shardsize: Number of samples in each shard. This in turn determines 
+                how much RAM is dedicated to hosting a shard in-memory when it is loaded.
+            shard_replace: If shards already exist, whether to replace them.
+            shardpath: Constructor for the directory where shards are kept. This
+                is appended with the `mode`. So, for training, a `shardpath` of 
+                "datashard-artifacts" becomes "datashard-artifacts-train"
+            shard-name: The name for each shard stored inside the `shardpath`.
+                This value is appended with the shard number. So a training shard
+                with index 5 and name "hf-shard" is stored at:
+                    "datashard-artifacts-train/hf-shard5.pt" 
+
+            maxlen: The maximum length of the text to encode. Usually 512.
+            mlm_probability: The probability of masking if using word and token masking.
+                Usually around 0.15
+            masking (bool): Whether to use any masking. Specific masking options are next. Can also be `train_only` for masking during training only.
+            token_mask: Whether to use random token masking, with `mlm_probability`
+            word_mask: Whether to use random word masking, with `mlm_probability`
+            keytoken_mask: Whether to use specific token masking, with `keytokens`
+            keyword_mask: Whether to use specific word masking, with `keywords`
+            keywords: What words to mask
+            keytokens: What tokens to mask
+            stopword_mask: Whether to use a stopword mask.
+            label_idxs: List of indices for labels. Defaults to [1].
+            annotation_idxs: List of indices for inputs. Defaults to [].
+
+
+            
+        """
+        self.dataset = dataset  # list of tuples or dict (text, *text, *labels)
+        self.logger = logger
+        self.file_len = len(dataset)
+        self.data_shuffle = kwargs.get("data_shuffle", True)
+
+        # Cacheing options
+        self.cache = kwargs.get("cache", False)
+        self.memcache = kwargs.get("memcache", False)
+        self.shardcache = kwargs.get("shardcache", False)
+
+        if int(self.cache) + int(self.memcache) + int(self.shardcache) > 1:
+            raise RuntimeError("Use only one cache type")
+
+        # Cache options
+        self.cache_replace = kwargs.get("cache_replace", False)
+        self.cachepath = kwargs.get("cachepath", "artifacts")
+        self.cachepath = mode+"-"+self.cachepath
+        self.cachename = kwargs.get("cachename", "h5-cache") + "-"  #the dash
+        if self.cache:
+            self.logger.debug("[Mode `{mode}`] Will look in path [{path}] for shards `{shards}[e].pt`".format(path=self.cachepath, shards=self.cachename, mode=mode))
+        self.base_cachepath = os.path.join(self.cachepath, self.cachename)
+        self.cache_exist = False
+        if os.path.exists(self.base_cachepath + ".h5"):
+            if self.cache_replace:
+                self.logger.debug("Deleting existing cache")
+                shutil.rmtree(self.cachepath)   # TODO this will delete the entire tree!!!!!!
+            else:
+                self.cache_exist = True
+                self.logger.debug("Cache already exists and `cache_replace` is False")
+        else:
+          self.logger.debug("Cache does not exist and will be created.")
+        if self.cache:
+            self.logger.debug("Creating cachepath %s"%self.cachepath)
+            os.makedirs(self.cachepath, exist_ok=True)
+
+
+        # Shardcache options
+        self.shardsize = kwargs.get("shardsize", 0)
+        self.shardsaveindex = 0
+        self.shard_replace = kwargs.get("shard_replace", False)
+        self.shardpath = kwargs.get("shardpath", "datashard-artifacts")
+        self.shardpath = mode+"-"+self.shardpath
+        self.shardname = kwargs.get("shardname", "fnc-filtermask-shard") + "-"  #the dash
+        if self.shardcache:
+            self.logger.debug("[Mode `{mode}`] Will look in path [{path}] for shards `{shards}[e].pt`".format(path=self.shardpath, shards=self.shardname, mode=mode))
+        self.base_shardpath = os.path.join(self.shardpath, self.shardname)
+        self.shards_exist = False
+        if os.path.exists(self.base_shardpath + "0.pt"):
+            if self.shard_replace:
+                self.logger.debug("Deleting existing shards")
+                shutil.rmtree(self.shardpath)
+            else:
+                self.shards_exist = True
+                self.logger.debug("Shards already exist and `shard_replace` is False")
+        else:
+          self.logger.debug("Shards do not exist and will be created.")
+        if self.shardcache:
+            self.logger.debug("Creating shardpath %s"%self.shardpath)
+            os.makedirs(self.shardpath, exist_ok=True)
+        
+        
+        
+        self.word_masking = kwargs.get("word_mask", False)
+        self.keyword_masking = kwargs.get("keyword_mask", True)
+        self.token_masking = kwargs.get("token_mask", False)
+        self.keytoken_masking = kwargs.get("keytoken_mask", False)
+        self.stopword_masking = kwargs.get("stopword_mask", False)
+
+        self.label_idxs = kwargs.get("label_idxs", [1])
+        self.annotation_idxs = kwargs.get("annotation_idxs", [])
+        
+        
+        # Tokenizer and lengths
+        self.tokenizer = kwargs.get("tokenizer")
+        self.maxlen = kwargs.get("maxlen")
+        self.mlm = kwargs.get("mlm_probability", 0.2)
+        self.masking = kwargs.get("masking", mode=="train")
+        if self.masking == "train_only":
+            self.masking = mode=="train"
+
+
+        self.keywords = []
+        if self.keyword_masking:
+            self.keywords = kwargs.get("keywords", [])
+        self.keytokens = []
+        if self.keytoken_masking:
+            self.keytokens = kwargs.get("keytokens", [])
+        # Store indices corresponding to keytokens
+        vocab = self.tokenizer.get_vocab()
+        self.keytokens = {vocab[item]:1 for item in self.keytokens if item in vocab}
+        self.keytokens_tensor = torch.tensor(list(self.keytokens.keys())).unsqueeze(0).T
+
+        # This is for memcache and shardcache
+        self.getcount = 0
+        self.shardcount = 0
+        if self.memcache or self.cache:
+            # Checks when we have processed entire dataset, refreshes mask ids when needed
+            self.refresh_flag = len(dataset)
+        if self.shardcache:
+            # We need to refresh mask ids each time we load a shard.
+            # This is ONLY for randomized masking, i.e. with MLM probability.
+            self.refresh_flag = self.shardsize
+            # this shuffles the internal shard index, so that we get a random example, instead of in order
+            self.shard_internal_shuffle = []
+            # This shuffles the shard index, so that each epoch, we load shards in different orders. We will set this up after setting up shards...
+            self.shard_shuffle = []
+        
+        self.getter = self.uncachedget  # for each entry, directly convert to features and return
+        if self.cache:
+            self.getter = self.diskget
+        elif self.memcache:
+            self.getter = self.memget
+        elif self.shardcache:
+            self.getter = self.shardget
+
+        if self.cache or self.memcache:
+        # The actual cache-ing
+            self.logger.debug("Started mem caching")
+            self.input_length_cache = []
+            self.logger.debug("Converting to features")
+            # For memcache, save dataset inside self.cached_dataset.
+            # Same for cache. Difference being, cached_dataset is either in memory or on disk!
+            self.convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen, cache=True)
+            self.logger.debug("Masking")
+            self.memcached_dataset = self.refresh_mask_ids(cache=True)
+
+        if self.shardcache:
+            self.input_length_cache = []
+            if not self.shards_exist:
+                self.logger.debug("Generating shards")
+                self.shardsaveindex = self.sharded_convert_to_features(self.dataset, self.tokenizer, maxlen=self.maxlen)    # save shards and get numshards
+            else:
+                self.shardsaveindex = len(glob(os.path.join(self.shardpath, "*.pt")))   # TODO Bug fix if files do not have consistent numbering
+            if self.shardsaveindex < 1:
+                raise ValueError("`shardsaveindex` is {val}, which is less than permissible minimum value of 1".format(val=self.shardsaveindex))
+            self.logger.debug("Obtained %i shards"%self.shardsaveindex)
+            self.shard_load_index = 0   # self.shardsaveindex is the maximum number of shards
+            self.shard_shuffle = list(range(self.shardsaveindex))    # count started from 0
+            if self.data_shuffle:
+                self.logger.debug("Shuffling shard load order")
+                random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+            self.current_shardsize = len(self.sharded_dataset)
+            if self.masking:
+                self.logger.debug("Refreshing token masks for loaded shard")    # TODO
+                self.sharded_dataset = self.sharded_refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+
+
+
+    def shardget(self, idx):
+        """Given an index, retrieve a single sample from the post-processed
+        dataset.
+
+        We retrieve from shards, so we need to track the current shard we are
+        working with, as well as which sample in the shard we are retrieving.
+
+        THus, given an index, we do not use that index directly. Instead,
+        HFGenerator tracks samples itself by using `self.shard_shuffle` 
+        and `self.internal_shard_shuffle`. The former tracks which shards
+        have been accessed already, and the latter tracks which samples
+        in the current shard have been tracked already.
+
+        `self.getcount` tracks the current index.
+
+        Shuffling if controlled with the `data_shuffle` argument during
+        initialization of the HFDataset.
+
+        Args:
+            idx (_type_): _description_
+
+        Returns:
+            _type_: _description_
+        """
+        # get entry from already loaded shardcache. so idx is incremented one by one -- No shuffling in data loader.
+        # This is a design choice, and we can't really do anything if user does shuffle + sharding
+        self.last_idx = idx
+        response = self.sharded_dataset[self.shard_internal_shuffle[self.getcount]]
+        self.getcount += 1
+
+        if self.getcount == self.current_shardsize:   # we have exhausted examples in this shard
+            self.getcount = 0
+            self.shard_load_index += 1                  # increment the shard index that we will load
+            if self.shard_load_index == self.shardsaveindex: # we have processed all shards.
+                self.shard_load_index = 0
+                if self.data_shuffle:
+                    random.shuffle(self.shard_shuffle)
+            self.sharded_dataset = self.load_shard(self.shard_shuffle[self.shard_load_index])
+
+            self.current_shardsize = len(self.sharded_dataset)
+            self.shard_internal_shuffle = list(range(self.current_shardsize))    #count started from 0
+            if self.data_shuffle:
+                random.shuffle(self.shard_internal_shuffle)
+            if self.masking:
+                self.sharded_dataset = self.sharded_refresh_mask_ids(self.sharded_dataset)  # TODO implement masking (and input_length_cache) for sharding
+
+        return response
+
+    def sharded_convert_to_features(self, dataset, tokenizer, maxlen):
+        features = []
+        shardsaveindex = 0
+        input_length_cache = []
+        input_word_length_cache = []
+        self.logger.info("Generating shards")
+        pbar = tqdm(total=int(self.file_len/self.shardsize)+1)
+        # TODO
+        # We basically need to deal with having multiple inputs and multiple labels...
+        # Ok, so, we already know the structure of the tuple that dataset gets
+        # Then, all we need is in config, specify additional_input_idx --> list | None
+        # if None, we don't need to do anything
+        # if a list, then our input is actually a tuple
+        # So, if a list, what do we do...?
+        # 
+        #
+        #
+        #
+        #
+        for idx, sample in enumerate(dataset):  # This is a tuple with (text, ..., stuff)
+            # Identify the indices of specific keywords to mask
+            #if self.keyword_masking or self.word_masking:
+            word_tokens = sample[0].split(" ")
+            encoded_word_length = len(word_tokens)
+            
+            if self.masking and self.keyword_masking:
+                keyword_idx = []
+                for widx,item in enumerate(word_tokens):
+                    if sum([1 if kword in item else 0 for kword in self.keywords]) > 0:
+                        keyword_idx.append(widx) # i.e. find word index for each keyword, if it exists
+
+            encoded = self.tokenizer(sample[0], return_tensors="pt", padding="max_length", max_length = maxlen, truncation = True, return_length = True)
+            encoded_token_length = torch.sum(encoded["attention_mask"])
+            encoded_word_ids = encoded.word_ids(0)
+            input_length_cache.append(encoded_token_length)
+            input_word_length_cache.append(encoded_word_length)
+            # create a list of lists. Each sublist is a mask for each keyword. So, we can AND all sublists to get the overall mask.
+            # Edge cases -- no keywords: 
+            merged_mask = encoded["attention_mask"]
+            if self.masking and self.keyword_masking:
+                match_idxs = torch.LongTensor([[word_idx != idx_of_keyword for word_idx in encoded_word_ids] for idx_of_keyword in keyword_idx])
+                if match_idxs.shape[0] > 0: #i.e. we have a mask for keywords, so we will and everything
+                    match_idxs = torch.all(match_idxs, dim=0, keepdim=True)
+                    merged_mask = torch.where(match_idxs == 0, match_idxs, merged_mask)
+
+            if self.masking and self.keytoken_masking:
+                # Dict with relevant keyword idxs...self.keytokens 
+                # First set 1 for all tokens, and 0 for masking tokens
+                match_idxs = encoded["input_ids"].repeat((len(self.keytokens),1))
+                match_idxs = torch.all(match_idxs != self.keytokens_tensor, dim=0, keepdim=True)
+                merged_mask = torch.where(match_idxs == 0, match_idxs, merged_mask)
+
+            
+            features.append(
+                (encoded["input_ids"], merged_mask, encoded["token_type_ids"], 
+                    encoded_token_length, encoded_word_length, encoded_word_ids, 
+                    [sample[label_idx] for label_idx in self.label_idxs],   # list of labels
+                    [sample[annot_idx] for annot_idx in self.annotation_idxs])   # list of secondary inputs. Defaults to empty.
+            )
+
+            if (idx+1)%self.shardsize == 0: # we need to save this shard.
+                # TODO need to handle multiple labels, later...
+                all_input_ids = torch.cat([f[0] for f in features])
+                # TODO
+                # Secondary inputs are other numeric inputs
+                # We check if secondary is provided by checking what is in  in features[7] (last element)
+                # if None, then secondary inputs are not provided. We can make it a matrix of zeros
+                # # if true, the secondary inputs are provided, in a list of numbers
+                # Then we convert then to a nxd matrix (n = shardsize, d = number of secondary inputs)
+                #all_secondary_inputs = None
+                all_attention_mask = torch.cat([f[1] for f in features])
+                all_token_type_ids = torch.cat([f[2] for f in features])
+                all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+                all_word_lens = torch.tensor([f[4] for f in features], dtype=torch.long)
+                all_word_ids = torch.tensor(np.nan_to_num(np.array([f[5] for f in features], dtype=float),nan=-1),dtype=torch.long)
+                all_labels = torch.tensor([f[6] for f in features], dtype=torch.long)   # TODO -- implement multiple labels as a list, then squash here to one dim if single label.
+                all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+                if len(features[0][7]) > 0:
+                    all_annotations = torch.tensor([f[7] for f in features])
+                else:
+                    all_annotations = torch.zeros((self.shardsize,1))
+
+                # Propagate input masking to output masking
+                for midx in range(all_attention_mask.shape[0]):
+                    masked_index = torch.where(all_attention_mask[midx][:input_length_cache[midx]]==0)[0]
+                    all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+                input_length_cache = []
+                input_word_length_cache = []
+
+                # SAVE HERE
+                self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations),
+                        shard_index = shardsaveindex)
+                shardsaveindex += 1
+                features = []
+                pbar.update(1)
+        # final shard...
+        if len(features) > 0:
+            all_input_ids = torch.cat([f[0] for f in features])
+            all_attention_mask = torch.cat([f[1] for f in features])
+            all_token_type_ids = torch.cat([f[2] for f in features])
+            all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+            all_word_lens = torch.tensor([f[4] for f in features], dtype=torch.long)
+            all_word_ids = torch.tensor(np.nan_to_num(np.array([f[5] for f in features], dtype=float),nan=-1),dtype=torch.long)
+            all_labels = torch.tensor([f[6] for f in features], dtype=torch.long)
+            all_masklm = -1*torch.ones(all_input_ids.shape, dtype=torch.long)
+            if len(features[0][7]) > 0:
+                    all_annotations = torch.tensor([f[7] for f in features])
+            else:
+                all_annotations = torch.zeros((self.shardsize,1))
+
+            # Propagate input masking to output masking
+            for midx in range(all_attention_mask.shape[0]):
+                masked_index = torch.where(all_attention_mask[midx][:input_length_cache[midx]]==0)[0]
+                all_masklm[midx][masked_index] = all_input_ids[midx][masked_index] # Set the masking labels for these to the actual word index from all_input_ids
+            input_length_cache = []
+
+            # SAVE HERE
+            self.save_shard(shard=TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations),
+                    shard_index = shardsaveindex)
+            shardsaveindex += 1
+            features = []
+            pbar.update(1)
+        pbar.close()
+        return shardsaveindex
+
+    def save_shard(self, shard: TensorDataset, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        torch.save(shard, shard_save_path)
+
+    def load_shard(self, shard_index: int):
+        shard_save_path = self.base_shardpath + str(shard_index) + ".pt"
+        self.logger.debug("Loading shard %s"%shard_save_path)
+        return torch.load(shard_save_path)
+
+    def memget(self, idx):
+        self.getcount+=1
+        if self.getcount > self.refresh_flag and self.masking:
+            self.memcached_dataset = self.refresh_mask_ids()
+            self.getcount = 0
+        return self.memcached_dataset[idx]
+
+    def refresh_mask_ids(self):
+        self.logger.debug("Refreshing mask ids")
+        if self.masking:
+            self.mask_ids = self.build_mask_ids(self.input_length_cache)
+            
+            all_attention_mask = self.all_attention_mask.clone()
+            all_masklm = self.all_masklm.clone()
+
+            for idx in range(self.refresh_flag):
+                all_attention_mask[idx][self.mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+                all_masklm[idx][self.mask_ids[idx]] = self.all_input_ids[idx][self.mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+
+            return TensorDataset(self.all_input_ids, all_attention_mask, self.all_token_type_ids, all_masklm, self.all_lens, self.all_labels)
+        else:
+            return TensorDataset(self.all_input_ids, self.all_attention_mask, self.all_token_type_ids, self.all_masklm, self.all_lens, self.all_labels)
+
+    def build_mask_ids(self, input_length_cache):
+        # for each element, we get a set of indices that are randomly selected...
+        # Also, -2 and +1 take care of [cls] and [sep] not being masked
+        return [(torch.randperm(inplength-2)+1)[:int(inplength*self.mlm)]  for inplength in input_length_cache]
+    
+    def build_whole_word_mask_ids(self, input_length_cache):
+        # for each element, we get a set of indices that are randomly selected...
+        # Also, -2 and +1 take care of [cls] and [sep] not being masked
+        return [(torch.randperm(word_length)+1)[:int(word_length*self.mlm)]  for word_length in input_length_cache]
+        
+
+        # So, given word_len, we can extract which ids need to be masked
+        # But, now we need a way to, given the ID that needs to be masked, get back the original words.
+    
+    def sharded_refresh_mask_ids(self, sharded_dataset: TensorDataset):
+        self.logger.debug("Refreshing mask ids")
+        #                   0               1                   2                 3           4           5             6            7............8
+        #TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations)
+                      
+        if self.token_masking or self.word_masking:
+            #merged_masklm =  torch.stack([shard[3] for shard in sharded_dataset])
+            all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+            all_attention_mask = torch.stack([shard[1] for shard in sharded_dataset])
+            if self.token_masking:
+                self.logger.debug("Performing random token masking")
+                mask_ids = self.build_mask_ids(torch.stack([shard[4] for shard in sharded_dataset]))
+                
+                
+                #all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+                all_input_ids = torch.stack([shard[0] for shard in sharded_dataset])
+
+                for idx in range(self.current_shardsize):
+                    all_attention_mask[idx][mask_ids[idx]] = 0 # Set the masking words to 0, so we do not attend to it during prediction
+                    all_masklm[idx][mask_ids[idx]] = all_input_ids[idx][mask_ids[idx]] # Set the masking labels for these to the actual word index from all_input_ids
+                #merged_masklm *= all_masklm
+            if self.word_masking:
+                self.logger.debug("Performing random word masking")
+                # TODO
+                masking_words = self.build_whole_word_mask_ids(torch.stack([shard[5] for shard in sharded_dataset]))
+                
+                #all_masklm = torch.stack([shard[3] for shard in sharded_dataset])
+                all_input_ids = torch.stack([shard[0] for shard in sharded_dataset])
+                all_word_ids = torch.stack([shard[6] for shard in sharded_dataset])
+                for idx in range(self.current_shardsize):
+                
+                    match_idxs = torch.LongTensor([[word_idx != idx_of_keyword for word_idx in all_word_ids[idx]] for idx_of_keyword in masking_words[idx]])
+                    if match_idxs.shape[0] > 0: #i.e. we have a mask for keywords, so we will and everything
+                        match_idxs = torch.all(match_idxs, dim=0, keepdim=True)
+                        merged_mask = torch.where(match_idxs[0]==0)[0]
+                        
+                        all_attention_mask[idx][merged_mask] = 0    # Set the masking words to 0, so we do not attend during prediction
+                        all_masklm[idx][merged_mask] = all_input_ids[idx][merged_mask] # Set the masking labels for these to the actual word index from all_input_ids
+                
+                #merged_masklm *= all_masklm
+            return TensorDataset(all_input_ids, all_attention_mask, torch.stack([shard[3] for shard in sharded_dataset]), all_masklm, 
+                torch.stack([shard[4] for shard in sharded_dataset]), 
+                torch.stack([shard[5] for shard in sharded_dataset]), 
+                torch.stack([shard[6] for shard in sharded_dataset]), 
+                torch.stack([shard[7] for shard in sharded_dataset]),
+                torch.stack([shard[8] for shard in sharded_dataset]))
+        else:
+            return sharded_dataset
+
+
+    def convert_to_features(self, dataset, tokenizer, maxlen, cache=False):
+        """Preprocess the dataset into text features. If cache is True, save the preprocessed data into a local h5 file.
+        Else, save in memory.
+
+        Args:
+            dataset (_type_): _description_
+            tokenizer (_type_): _description_
+            maxlen (_type_): _description_
+            cache (bool, optional): _description_. Defaults to False.
+        """
+        features = []
+        self.input_length_cache = []
+        for idx, sample in enumerate(dataset):
+            tokens = self.tokenizer.tokenize(sample[0])
+            if len(tokens) > maxlen - 2:
+                tokens = tokens[0:(maxlen - 2)]
+            # An easy way potentially -- get the index of each word in the token set...
+            # Then check which word is in keyword
+            # Then mask out that word...????????
+            finaltokens = ["[CLS]"]
+            token_type_ids = [0]
+            for token in tokens:
+                finaltokens.append(token)
+                token_type_ids.append(0)
+            finaltokens.append("[SEP]")
+            token_type_ids.append(0)
+
+            input_ids = self.tokenizer.convert_tokens_to_ids(finaltokens)
+            attention_mask = [1]*len(input_ids)
+            input_len = len(input_ids)
+            self.input_length_cache.append(len(input_ids))
+            while len(input_ids) < maxlen:
+                input_ids.append(0)
+                attention_mask.append(0)
+                token_type_ids.append(0) 
+
+            assert len(input_ids) == maxlen
+            assert len(attention_mask) == maxlen
+            assert len(token_type_ids) == maxlen
+        
+            features.append(
+                (input_ids, attention_mask, token_type_ids, input_len, *sample[1:])
+            )
+
+        self.all_input_ids = torch.tensor([f[0] for f in features], dtype=torch.long)
+        self.all_attention_mask = torch.tensor([f[1] for f in features], dtype=torch.long)
+        self.all_token_type_ids = torch.tensor([f[2] for f in features], dtype=torch.long)
+        self.all_lens = torch.tensor([f[3] for f in features], dtype=torch.long)
+        self.all_labels = torch.tensor([f[4] for f in features], dtype=torch.long)
+        self.all_masklm = -1*torch.ones(self.all_input_ids.shape, dtype=torch.long)
+                
+    def __len__(self):
+        return len(self.dataset)
+
+    def __getitem__(self, idx):
+        return self.getter(idx)
+
+    def uncachedget(self, idx):
+        return 
+
+
+
+
+
+from ednaml.utils.LabelMetadata import LabelMetadata
+from ednaml.generators import TextGenerator
+class HFGeneratorOld(TextGenerator):
+    """HFGenerator is a generic text generator for HuggingFace-tokenizers
+    and designed to work with HuggingFace style transformers.
+
+    It can use any tokenizer that follows HuggingFace API.
+
+    Usage:
+        TODO
+    """
+
+    def build_transforms(self, transform, mode, **kwargs):  #<-- generator kwargs:
+        """Use the GENERATOR_KWARGS to identify the tokenizer
+        class. For now, only built-in tokenizers are supported.
+
+        Args:
+            transform (_type_): Unused. It is necessary here as a dummy argument in 
+                constructor for the core ednaml.generator API.
+            mode (_type_): Whether to use train or test mode. Useful when selecting
+                the correct subset from crawler.
+        """
+        from ednaml.utils import locate_class
+        print("Building Transforms")
+        tokenizer = kwargs.get("tokenizer", "HFAutoTokenizer")
+        self.tokenizer = locate_class(package="ednaml", subpackage="utils", classpackage=tokenizer, classfile="tokenizers")
+        self.tokenizer = self.tokenizer(**kwargs) # vocab_file, do_lower_case, spm_model_file
+
+    
+    def buildDataset(self, crawler, mode, transform, **kwargs): #<-- dataset args:
+        print("Building Dataset")
+        return HFDataset(self.logger, crawler.metadata[mode]["crawl"], mode, tokenizer = self.tokenizer, **kwargs) # needs maxlen, memcache, mlm_probability
+
+    def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
+        print("Building Dataloader")
+        return torch.utils.data.DataLoader(dataset, batch_size=batch_size*self.gpus,
+                                            shuffle=kwargs.get("shuffle", False), num_workers = self.workers, 
+                                        collate_fn=self.collate_fn)
+
+    def getNumEntities(self, crawler, mode, **kwargs):  #<-- dataset args
+        label_dict = {
+            item: {"classes": crawler.metadata[mode]["classes"][item]}
+            for item in kwargs.get("classificationclass", ["color"])
+        }
+        return LabelMetadata(label_dict=label_dict)
+
+    def collate_fn(self, batch):
+        all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_lens, all_word_lens, all_word_ids, all_labels, all_annotations = map(torch.stack, zip(*batch))
+        max_len = max(all_lens).item()
+        all_input_ids = all_input_ids[:, :max_len]
+        all_attention_mask = all_attention_mask[:, :max_len]
+        all_token_type_ids = all_token_type_ids[:, :max_len]
+        all_masklm = all_masklm[:, :max_len]
+        # Right now, we don't care that much about all_word_ids or all_word_lens
+
+        return all_input_ids, all_attention_mask, all_token_type_ids, all_masklm, all_annotations, all_labels
+
+
+
+
diff --git a/src/ednaml/generators/ImageGenerator.py b/src/ednaml/generators/ImageGenerator.py
index 6a201f44..1bc1e19c 100644
--- a/src/ednaml/generators/ImageGenerator.py
+++ b/src/ednaml/generators/ImageGenerator.py
@@ -58,17 +58,17 @@ def _build_transforms(
         )
         transformer_primitive = []
         transformer_primitive.append(T.Resize(size=i_shape))
-        if kwargs.get("h_flip") > 0:
+        if kwargs.get("h_flip", 0) > 0:
             transformer_primitive.append(
                 T.RandomHorizontalFlip(p=kwargs.get("h_flip"))
             )
-        if kwargs.get("t_crop"):
+        if kwargs.get("t_crop", False):
             transformer_primitive.append(T.RandomCrop(size=i_shape))
         transformer_primitive.append(T.ToTensor())
         transformer_primitive.append(
             T.Normalize(mean=normalization_mean, std=normalization_std)
         )
-        if kwargs.get("rea"):
+        if kwargs.get("rea", False):
             transformer_primitive.append(
                 T.RandomErasing(
                     p=0.5, scale=(0.02, 0.4), value=kwargs.get("rea_value", 0)
diff --git a/src/ednaml/generators/MultiClassificationGenerator.py b/src/ednaml/generators/MultiClassificationGenerator.py
index 4735069b..88365eac 100644
--- a/src/ednaml/generators/MultiClassificationGenerator.py
+++ b/src/ednaml/generators/MultiClassificationGenerator.py
@@ -18,7 +18,7 @@ def __init__(self, dataset, transform=None, **kwargs):
         self.transform = transform
         self.pathidx = kwargs.get("pathidx", 0)
         self.annotationidx = kwargs.get(
-            "annotationidx", 1
+            "annotationidx", [1]
         )  # This should be a list...
         self.num_annotations = len(self.annotationidx)
 
diff --git a/src/ednaml/generators/TorchvisionGeneratorWrapper.py b/src/ednaml/generators/TorchvisionGeneratorWrapper.py
index 1d381c4c..6ab3288f 100644
--- a/src/ednaml/generators/TorchvisionGeneratorWrapper.py
+++ b/src/ednaml/generators/TorchvisionGeneratorWrapper.py
@@ -4,7 +4,8 @@
 from torch.utils.data import DataLoader as TorchDataLoader
 from ednaml.utils import locate_class
 from torchvision.datasets import VisionDataset
-
+import torch
+from ednaml.utils.LabelMetadata import LabelMetadata
 
 class TorchvisionGeneratorWrapper(ImageGenerator):
     """Generator for a TorchVision dataset.
@@ -13,35 +14,12 @@ class TorchvisionGeneratorWrapper(ImageGenerator):
         ImageGenerator (_type_): _description_
     """
 
-    def __init__(
-        self,
-        logger, 
-        gpus: int,
-        i_shape: Union[List[int], Tuple[int, int]],
-        channels: int,
-        normalization_mean: float,
-        normalization_std: float,
-        normalization_scale: float,
-        **kwargs
-    ):
-        super().__init__(logger,
-            gpus,
-            i_shape,
-            channels,
-            normalization_mean,
-            normalization_std,
-            normalization_scale,
-            **kwargs
-        )
-
-        """
-        So we will have arguments inside the kwargs that set up the dataloader object
-        """
-        self.gpus = gpus
+    def buildGeneratorAttributes(self, **kwargs):
         self.torchvision_dataset_class = kwargs.get("tv_dataset")
         self.torchvision_dataset_args = kwargs.get("tv_args")
 
-    def build(self, datacrawler, mode, batch_size, workers, **kwargs):
+        
+    def buildDataset(self, datacrawler, mode: str, transform: List[object], **kwargs) -> TorchDataset:
         # We don't need the data crawler...
 
         dataset_class: VisionDataset = locate_class(
@@ -54,14 +32,14 @@ def build(self, datacrawler, mode, batch_size, workers, **kwargs):
         # dataset API is not unified, see discussions below
         # https://github.com/pytorch/vision/issues/1067
         # https://github.com/pytorch/vision/issues/1080
-        if self.torchvision_dataset_class in ["CIFAR10", "MNIST", "USPS"]:
+        if self.torchvision_dataset_class in ["CIFAR10", "MNIST", "USPS", "CIFAR100"]:
             # These have a train argument...
-            self.dataset = dataset_class(
+            return dataset_class(
                 root=self.torchvision_dataset_args.get("root"),
                 train=(mode == "train"),
                 transform=self.transformer,
                 target_transform=None,
-                **self.torchvision_dataset_args
+                **self.torchvision_dataset_args.get("args", {})
             )
         elif self.torchvision_dataset_class in [
             "CelebA",
@@ -72,19 +50,16 @@ def build(self, datacrawler, mode, batch_size, workers, **kwargs):
         ]:
             if self.torchvision_dataset_class == "ImageNet":
                 mode = mode if mode == "train" else "val"
-            self.dataset = dataset_class(
+            return dataset_class(
                 root=self.torchvision_dataset_args.get("root"),
                 split=mode,
                 transform=self.transformer,
                 target_transform=None,
-                **self.torchvision_dataset_args
+                **self.torchvision_dataset_args.get("args", {})
             )
         else:
             raise NotImplementedError()
 
-        self.dataloader = self.buildDataLoader(
-            self.dataset, mode, batch_size=batch_size, **kwargs
-        )
 
     def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
         return TorchDataLoader(
@@ -93,3 +68,10 @@ def buildDataLoader(self, dataset, mode, batch_size, **kwargs):
             shuffle=kwargs.get("shuffle", True),
             num_workers=self.workers,
         )
+
+    def getNumEntities(self, datacrawler, mode, **kwargs) -> LabelMetadata:
+        label_name = kwargs.get("label_name", "label")
+        num_classes = kwargs.get("num_classes", None)
+        if num_classes is None:
+            num_classes = torch.unique(torch.Tensor(self.dataset.targets)).shape[0]
+        return LabelMetadata({label_name: {"classes": num_classes}})
\ No newline at end of file
diff --git a/src/ednaml/generators/__init__.py b/src/ednaml/generators/__init__.py
index 8b675ac0..0c98cd3f 100644
--- a/src/ednaml/generators/__init__.py
+++ b/src/ednaml/generators/__init__.py
@@ -50,7 +50,12 @@ def __init__(
         self.transformer = self.build_transforms(
             self.transforms, self.mode, **kwargs
         )
-
+        self.buildGeneratorAttributes(**kwargs)	
+    
+    
+    def buildGeneratorAttributes(self, **kwargs):	
+        pass
+    
     def build_transforms(self, transforms: Dict[str, Any], mode, **kwargs):
         return None
 
@@ -125,3 +130,4 @@ def getNumEntities(self, datacrawler, mode, **kwargs) -> LabelMetadata:
 from ednaml.generators.TorchvisionGeneratorWrapper import (
     TorchvisionGeneratorWrapper,
 )
+from ednaml.generators.HFGenerator import HFGenerator
diff --git a/src/ednaml/models/HFAutoModel.py b/src/ednaml/models/HFAutoModel.py
new file mode 100644
index 00000000..755d59d1
--- /dev/null
+++ b/src/ednaml/models/HFAutoModel.py
@@ -0,0 +1,58 @@
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.utils import locate_class
+from transformers import AutoConfig
+
+class HFAutoModel(ModelAbstract):
+    """For Sequence Classification only.
+
+    Future plans: HF has a pretty structured method for the heads
+    So, we can extend the HFAutoModel to use multiple heads...
+
+    Essentially, construct the necessary heads from the base model type 
+    (i.e. <ModelName>ForSequenceClassification, i.e. MPNetForSequenceClassification)
+
+    Then, we construct the core model, i.e <ModelName>Model.from_pretrained(), i.e. MPNetModel.from_pretrained()
+    Then, we can stick the heads on top.
+
+    Args:
+        ModelAbstract (_type_): _description_
+
+    MODEL_KWARGS:
+        auto_class: Which of the HuggingFace AutoModels to use. Choose from:
+            AutoModel, AutoModelForPreTraining, AutoModelWithMLMHead, AutoModelForSequenceClassification, AutoModelForQuestionAnswering, AutoModelForTokenClassification
+        from_pretrained: The base for the model. If not provided, will default to value in MODEL_BASE
+        **kwargs: Any other HF kwargs for an AutoModel
+
+    Model_Kwargs:
+        NOTE: Since we are loading from pretrained, either MODEL_BASE or MODEL_KWARGS.from_pretrained MUST be passed
+        NOTE: Do NOT use `pretrained_model_name_or_path` as a MODEL_KWARGS field, because this will conflict with passed arguments to AutoModel. 
+        See HuggingFace API for details on model_kwargs
+
+    """
+
+
+    def model_attributes_setup(self, **kwargs):
+        self.from_pretrained = kwargs.get("from_pretrained", self.model_base)
+        self.auto_class = kwargs.get("auto_class", "AutoModel")
+
+
+    def model_setup(self, **kwargs):
+        auto_class = locate_class(package="transformers", subpackage=self.auto_class)
+        self.autoconfig = AutoConfig.from_pretrained(self.from_pretrained, **kwargs)
+        self.encoder = auto_class.from_pretrained(self.from_pretrained, config=self.autoconfig)
+
+    def forward_impl(self, x, 
+                        attention_mask=None,
+                        token_type_ids=None,
+                        position_ids=None,
+                        head_mask=None,
+                        inputs_embeds=None,
+                        labels=None,
+                        output_attentions=None,
+                        output_hidden_states=None, **kwargs):
+        response = self.encoder(x,attention_mask = attention_mask, token_type_ids = token_type_ids, position_ids=position_ids, head_mask=head_mask, 
+                                    inputs_embeds = inputs_embeds, labels=labels, output_attentions=output_attentions, output_hidden_states=output_hidden_states)
+                                    
+        # https://stackoverflow.com/questions/61323621/how-to-understand-hidden-states-of-the-returns-in-bertmodelhuggingface-transfo
+        return response.logits, response.hidden_states[-1][:,0,:], [response.attentions]
+        
\ No newline at end of file
diff --git a/src/ednaml/models/ModelAbstract.py b/src/ednaml/models/ModelAbstract.py
index 268885c5..8014b7b7 100644
--- a/src/ednaml/models/ModelAbstract.py
+++ b/src/ednaml/models/ModelAbstract.py
@@ -1,16 +1,39 @@
-from torch import nn
+from logging import Logger
+from os import PathLike
+from torch import TensorType, nn
 import torch
-from typing import List, Dict
+from typing import Any, List, Dict, Tuple, Type
+from ednaml.plugins import ModelPlugin
 from ednaml.utils.LabelMetadata import LabelMetadata
 
 
 class ModelAbstract(nn.Module):
-    model_name = "ModelAbstract"
-    model_arch = None
-    number_outputs = 1
-    output_classnames = ["out1"]
-    output_dimensions = [512]
-    secondary_outputs = []
+    """ModelAbstract is the base class for an Edna Model used in EdnaML, EdnaDeploy, and any other Edna frameworks. ModelAbstract allows design of practically any neural-network based trainable model by inheriting from torch.nn.
+
+    ModelAbstract allows users to define a model by following specific recipes. While reducing some flexibibility, the end result is to allow for better interoperability and extensibility.
+
+    Attributes:
+        model_name (str): The name for this model. Useful during debugging.
+        model_arch (TODO): TODO
+        number_outputs (int): The number of outputs for this model
+        output_classnames (List[str]): The name for each output for this model. This can be set manually or programatically.
+        output_dimensions (List[int]): The dimensions for each output for this model. This can be set manually or programatically.
+        secondary_outputs (TODO): TODO
+        model_base (str): Name for the core model architecture, of this is a modular construction. Useful during debugging.
+        weights (str): Path to weights file for this model
+        normalization (str): Name for type of normalization used in this model. Provided here since it is a common model characteristic.
+        metadata (LabelMetadata): A LabelMetadata object for this model that provides information on the training input and mopdel output. 
+        parameter_groups (Dict[str, nn.Module]): A named dictionary of a ModelAbstracts sub-components. For example, a GAN can have encoder, decoder, and discriminator parameter groups, if they are trained separately.
+
+    Methods:
+        _type_: _description_
+    """
+    model_name: str = "ModelAbstract"
+    model_arch: str = None
+    number_outputs: int = 1
+    output_classnames: List[str] = ["out1"]
+    output_dimensions: List[int] = [512]
+    secondary_outputs: List[Any] = []
 
     model_base: str
     weights: str
@@ -18,10 +41,17 @@ class ModelAbstract(nn.Module):
     metadata: LabelMetadata
     parameter_groups: Dict[str, nn.Module]
 
+    plugins: Dict[str,ModelPlugin] = {}
+    plugin_count: int = 0
+    has_plugins: bool = False
+    plugin_hook: str = "always"
+
+    _logger: Logger = None
+
     def __init__(
         self,
-        base=None,
-        weights=None,
+        base: str = None,
+        weights: PathLike = None,
         metadata: LabelMetadata = None,
         normalization: str = None,
         parameter_groups: List[str] = None,
@@ -34,11 +64,15 @@ def __init__(
         self.normalization = normalization
         self.parameter_groups = {}
         self.inferencing = False
+        
 
         self.model_attributes_setup(**kwargs)
         self.model_setup(**kwargs)
         self.parameter_groups_setup(parameter_groups)
 
+    def set_plugin_hooks(self, plugin_hook: str):   # always | warmup | activated
+        self.plugin_hook = plugin_hook
+
     def model_attributes_setup(self, **kwargs):
         raise NotImplementedError()
 
@@ -93,19 +127,28 @@ def partial_load(self, weights_path):
                 continue
             self.state_dict()[_key].copy_(params[_key])
 
-    def forward(self, x, **kwargs):
+    def forward(self, x, labels=None, **kwargs):    # TODO Labels for the plugins...?????
         if self.training and self.inferencing:
             raise ValueError(
                 "Cannot inference and train at the same time! Call"
                 " deinference() first, before train()"
             )
+
+        # TODO, we neede a pre-forward hook, a forward hook, and a post foward hook
+        if self.has_plugins:
+            x, kwargs, secondary_output_queue_pre = self.pre_forward_hook(x, **kwargs)
         feature_logits, features, secondary_outputs = self.forward_impl(
             x, **kwargs
         )
+        if self.has_plugins:
+            feature_logits, features, secondary_outputs, secondary_output_queue_post = self.post_forward_hook(x, feature_logits, features, secondary_outputs, **kwargs)
 
+        # TODO deal with secondary_output_queues
+        if self.has_plugins:
+            secondary_outputs = (secondary_outputs, secondary_output_queue_pre, secondary_output_queue_post)
         return feature_logits, features, secondary_outputs
 
-    def foward_impl(self, x, **kwargs):
+    def foward_impl(self, x, **kwargs) -> Tuple[TensorType,TensorType,List[Any]]:   # Return Logits, Features, Secondary Outputs
 
         raise NotImplementedError()
 
@@ -145,3 +188,85 @@ def eval(self):
 
     def convertForInference(self) -> "ModelAbstract":
         raise NotImplementedError
+
+    #-------------------------------------------------------------------------------------------
+    # Model Plugins and Hooks architecture
+    def loadPlugins(self, plugin_path: PathLike, ignore_plugins: List[str] = []):
+        plugin_dict = torch.load(plugin_path)
+        for plugin_name in plugin_dict:
+            if plugin_name in self.plugins:
+                if plugin_name not in ignore_plugins:
+                    self.plugins[plugin_name].load(plugin_dict[plugin_name])
+                    self._logger.info("Loading plugin with name %s"%plugin_name)
+                else:
+                    self._logger.info("Ignoring plugin with name %s"%plugin_name)
+            else:
+                self._logger.info("No plugin exists for name %s"%plugin_name)
+
+    def savePlugins(self):
+        save_dict = {}
+        for plugin_name in self.plugins:
+            save_dict[plugin_name] = self.plugins[plugin_name].save()
+            
+
+        return save_dict
+
+    def addPlugin(self, plugin: Type[ModelPlugin], plugin_name: str = None, plugin_kwargs: Dict[str, Any] = {}):
+        if plugin_name is None:
+            plugin_name = plugin.name
+        if plugin_name == "ModelPlugin":
+            # TODO NOTE we are badly changing the ModelPlugin name to the class name
+            raise ValueError("Potentially no actual plugin passed!")
+
+        if plugin_name in self.plugins:
+            raise KeyError("`plugin_name` %s already exists in self.plugins:  "%plugin_name)
+        else:
+            self._logger.info("Added plugin %s"%plugin_name)
+            self.plugins[plugin_name] = plugin(**plugin_kwargs)
+            self.plugins[plugin_name]._logger = self._logger
+        
+        self.plugin_count = len(self.plugins)
+        self.has_plugins = self.plugin_count > 0
+
+    def set_usable_plugins(self):
+        if self.plugin_hook == "always":
+            self.usable_plugins = [item for item in self.plugins]
+        elif self.plugin_hook == "warmup":
+            self.usable_plugins = [item for item in self.plugins if not self.plugins[item].activated]
+        elif self.plugin_hook == "activated":
+            self.usable_plugins = [item for item in self.plugins if self.plugins[item].activated]
+        else:
+            raise ValueError("Unknown value for `plugin_hook`: %s"%self.plugin_hook)
+
+    def pre_epoch_hook(self, epoch: int = 0):
+        # Here, we set up which plugins are activated and can be used, based on cfg...PLUGIN.HOOK
+        # we do this check before, and after...
+        self.set_usable_plugins()   # dealing with a plugin-load
+        for plugin in self.usable_plugins:
+            self.plugins[plugin].pre_epoch(model = self, epoch=epoch)
+        self.set_usable_plugins()   # post plugin usage
+    def post_epoch_hook(self, epoch: int = 0):
+        # Here, we set up which plugins are activated and can be used, based on cfg...PLUGIN.HOOK
+        # we do this check before, and after...
+        #self.set_usable_plugins()  # unneeded
+        for plugin in self.usable_plugins:
+            self.plugins[plugin].post_epoch(model = self, epoch=epoch)
+        self.set_usable_plugins()   # post plugin usage
+
+    def pre_forward_hook(self, x, **kwargs):
+        # For example, L-Score will perturb the input, then itself call self.forward_impl with the perturbed input (i.e. itself very 
+        # hacky under our own framework), then provide results to secondary_outputs
+        secondaries: Dict[str, Any] = {}
+        for plugin in self.usable_plugins:
+            x, kwargs, secondary_output_pre = self.plugins[plugin].pre_forward(x, **kwargs)
+            secondaries[plugin] = secondary_output_pre
+        return x, kwargs, secondaries
+
+    def post_forward_hook(self, x, feature_logits, features, secondary_outputs, **kwargs):
+        # For ModelPlugins.
+        # For example, KMP, after forward pass, provides distance to nearest proxy as well as nearest proxy in the secondary_outputs
+        secondaries: Dict[str, Any] = {}
+        for plugin in self.usable_plugins:
+            feature_logits, features, secondary_outputs, kwargs, secondary_output_post = self.plugins[plugin].post_forward(x,feature_logits, features, secondary_outputs,self, **kwargs)
+            secondaries[plugin] = secondary_output_post
+        return feature_logits, features, secondary_outputs, secondaries
diff --git a/src/ednaml/models/MultiBranchResnet.py b/src/ednaml/models/MultiBranchResnet.py
index 76637654..94f6c8c3 100644
--- a/src/ednaml/models/MultiBranchResnet.py
+++ b/src/ednaml/models/MultiBranchResnet.py
@@ -410,12 +410,12 @@ def forward_impl(self, x, **kwargs):
             fused_outs += [self.softmax_fused(fused_features[0])]
 
         # in forward_impl, we do for softtargetnames in self.soft-t-b-n, self.branches[softtargetnames].features --> send to the thingamagig and add to outputs
-        soft_outs = []
+        soft_outs = [None]*len(self.soft_target_branch_names)
         for idx, softtargetnames in enumerate(self.soft_target_branch_names):
             soft_outs[idx] = self.soft_target_outputs[idx](
                 features[self.branch_name_idx_map[softtargetnames]]
             )
-        return features + fused_features, outputs + fused_outs + soft_outs, []
+        return outputs + fused_outs + soft_outs, features + fused_features, []
 
     def parameter_groups_setup(self, parameter_groups: List[str]):
         self.parameter_groups[parameter_groups[0]] = self
diff --git a/src/ednaml/optimizer/__init__.py b/src/ednaml/optimizer/__init__.py
index ea8fbb0f..ddcd614a 100644
--- a/src/ednaml/optimizer/__init__.py
+++ b/src/ednaml/optimizer/__init__.py
@@ -62,7 +62,7 @@ def build(self, model: ModelAbstract) -> torch.optim.Optimizer:
                     learning_rate = self.base_lr * self.lr_bias
                     weight_decay = self.weight_decay * self.weight_bias
                 else:
-                    learning_rate = self.base_lr * self.gpus
+                    learning_rate = self.base_lr * max(1,self.gpus)
                     weight_decay = self.weight_decay
                 params += [
                     {
diff --git a/src/ednaml/plugins/FastKMeansProxy.py b/src/ednaml/plugins/FastKMeansProxy.py
new file mode 100644
index 00000000..62ee7719
--- /dev/null
+++ b/src/ednaml/plugins/FastKMeansProxy.py
@@ -0,0 +1,175 @@
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+from sklearn.cluster import MiniBatchKMeans
+import h5py, os
+from ednaml.plugins.ModelPlugin import ModelPlugin
+
+class FastKMeansProxy(ModelPlugin):
+    name = "FastKMeansProxy"
+    def __init__(self, proxies = 10, dimensions = 768, dist = "euclidean", rand_seed = 34623498, iterations = 25, batch_size = 256, alpha = 0.5, 
+                        feature_file = None, classifier_access = "classifier", **kwargs):
+        super().__init__(proxies = proxies, dimensions = dimensions, dist = dist, rand_seed = rand_seed, iterations = iterations, batch_size = batch_size, 
+                        alpha = alpha, feature_file = feature_file, classifier_access = classifier_access)
+
+
+    def build_plugin(self, **kwargs):
+        self.proxies = kwargs.get("proxies")
+        self.rand_seed = kwargs.get("rand_seed")
+        self.dimensions = kwargs.get("dimensions")
+        self.dist = kwargs.get("dist")    # TODO
+        self.batch_size = kwargs.get("batch_size")
+        self.iterations = kwargs.get("iterations")
+        self.alpha = kwargs.get("alpha")
+
+        self.activated = False
+        self._dist_func = self.build_dist(self.dist)
+        if self.dist == "euclidean":
+            self._preprocess = lambda x:x
+        elif self.dist == "cosine":
+            self._preprocess = torch.nn.functional.normalize
+
+        self.classifier_access = kwargs.get("classifier_access")
+        self._classifier_setup = False
+        self._classifier = None
+
+
+        output_file = kwargs.get("feature_file")
+        if output_file is None:
+            import uuid
+            output_file = "-".join(["FKMP", "features", uuid.uuid1().urn.split("-")[0].split(":")[-1]])
+        self._output_file = output_file + ".h5"
+        if os.path.exists(self._output_file):
+            os.remove(self._output_file)
+        self._writer = h5py.File(self._output_file, "w")  #we will delete old training features file
+        self._written = False
+        self._prev_idx = -1
+
+
+    def build_dist(self, dist="euclidean"):
+        if dist == "euclidean":
+            return self.l2dist
+        elif dist == "cosine":
+            return self.cosdist
+        else:
+            raise ValueError("Invalid value for dist: %s"%dist)
+    def l2dist(self, x):
+        return torch.argmin(torch.sqrt(((self.cluster_means - x)**2).sum(1)))
+    def cosdist(self, x):   #https://stackoverflow.com/questions/46409846/using-k-means-with-cosine-similarity-python
+        return self.l2dist(torch.nn.functional.normalize(x, dim=0))
+
+    def build_params(self, **kwargs):
+        self.cluster_means = torch.rand(self.proxies, self.dimensions)
+        self.inertia = 0
+        self.cluster_counts = torch.zeros(self.proxies)
+        self.kdcluster = None
+        self.high_density_thresholds = []
+        self.proxy_label = []
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self._classifier_setup:
+            self._classifier = {mname:mlayer for mname, mlayer in model.named_modules()}[self.classifier_access]
+            self._classifier_setup = True
+        if not self.activated:
+            # perform the training here
+            self.save_features(features)
+            
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            dist, labels, proxy_labels = self.compute_labels(features)
+            return feature_logits, features, secondary_outputs, kwargs, {"threshold": labels, "distance": dist, "proxy_labels": proxy_labels}
+
+    def save_features(self, features):
+        feats = features.cpu().numpy()
+        if self._written:
+            self._writer["features"].resize((self._writer["features"].shape[0] + feats.shape[0]), axis=0)
+            self._writer["features"][-feats.shape[0]:] = feats        
+        else:   # First time writing -- we will need to create the dataset.
+            self._writer.create_dataset("features", data=feats, compression = "gzip", chunks=True, maxshape=(None,feats.shape[1]))
+            self._written = True
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        try:    # in case it is closed already somehow...
+            self._writer.close()
+        except:
+            pass
+        if not self.activated:  # If already activated, we do not need to change anything
+            self.post_epoch_num = epoch    # Maybe have a better check in case something drops or epochs get skipped...
+            if self.pre_epoch_flag: # Means we had a pre_epoch flag so an epoch was completed.
+                self.pre_epoch_flag = False # reset the flag
+                if self.post_epoch_num != self.pre_epoch_num:
+                    raise ValueError("Epoch may have been skipped. Before: %i\tAfter: %i"%(self.pre_epoch_num, self.post_epoch_num))
+            
+            self.performkmeans()
+            self.highdensitybins(model)
+            with torch.no_grad():
+                self.proxy_label = torch.argmax(self._classifier(self.cluster_means.cuda()).cpu(), dim=1)
+            self.high_density_thresholds  = torch.tensor(self.high_density_thresholds)
+            self.activated = True
+
+    def performkmeans(self):
+        import time
+        inertia = []
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        
+        print("Starting KMeans for k={kval}, with {iters} iterations".format(kval=self.proxies, iters=self.iterations))
+        kmeans = MiniBatchKMeans(n_clusters = self.proxies, random_state = self.rand_seed, batch_size = self.batch_size)
+        stime = time.time()
+        for iters in range(self.iterations):
+            for i in range(0, data_size, self.batch_size):
+                current_data = data['features'][i:i+self.batch_size]
+                kmeans.partial_fit(current_data)
+            if iters%5 == 0:
+                etime = round(time.time() - stime, 2)
+                print("\t[{elapse} s] -- Completed {iters} iterations".format(iters=iters, elapse = etime))
+                stime = time.time()
+        print("\tCompeted MBKM for k={kval}, with inertia: {inertia}".format(kval=self.proxies, inertia = kmeans.inertia_))
+        data.close()
+        self.cluster_means = torch.tensor(kmeans.cluster_centers_)
+        self.inertia = kmeans.inertia_
+        return inertia
+
+    def highdensitybins(self, model: ModelAbstract):
+        if self.dist == "euclidean":
+            self.kdcluster = KDTree(self.cluster_means)
+        elif self.dist == "cosine":
+            self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+        else:
+            raise ValueError("Invalid value for dist: %s"%self.dist)
+
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        distance_bins = [[] for _ in range(self.proxies)]
+        print("Starting High Density estimation")
+        for i in range(0, data_size, self.batch_size):
+            current_data = data['features'][i:i+self.batch_size]
+            dist, indices = self.kdcluster.query(self._preprocess(torch.tensor(current_data)), k=1, return_distance=True)   #.squeeze()
+            for idx,val in enumerate(indices[:,0]):
+                distance_bins[val].append(dist[idx, 0])
+        
+        self.high_density_thresholds = [None]*self.proxies
+        import numpy as np
+        for proxy in range(self.proxies):
+            self.high_density_thresholds[proxy] = np.percentile(distance_bins[proxy], self.alpha * 100)
+        print("Completed High Density threshold estimation")
+        data.close()
+
+    def compute_labels(self, features):
+        """Compute the cluster labels for dataset X given centers C.
+        """
+        # labels = np.argmin(pairwise_distances(C, X), axis=0) # THIS REQUIRES TOO MUCH MEMORY FOR LARGE X
+        feats = features.cpu()
+        dist, idx = self.kdcluster.query(self._preprocess(feats), k=1, return_distance=True)   #.squeeze()
+        # TODO convert idx to the actual cluster means to the actual cluster labels...
+        # Returns distance  | high-density-threshold    | index inside cluster_means. We will need to look up label later...
+        #return torch.from_numpy(dist).squeeze(1), torch.tensor([self.high_density_thresholds[item[0]] for item in idx]), idx[:,0]
+        return torch.from_numpy(dist).squeeze(1), self.high_density_thresholds[idx[:,0]], self.proxy_label[idx[:,0]]
+
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
diff --git a/src/ednaml/plugins/FastRandomizedLipschitz.py b/src/ednaml/plugins/FastRandomizedLipschitz.py
new file mode 100644
index 00000000..ea81c104
--- /dev/null
+++ b/src/ednaml/plugins/FastRandomizedLipschitz.py
@@ -0,0 +1,318 @@
+
+from typing import List
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+from ednaml.plugins.ModelPlugin import ModelPlugin
+from sortedcontainers import SortedKeyList
+import numpy as np, os, h5py
+from sklearn.cluster import MiniBatchKMeans
+
+
+class FastRandomizedLipschitz(ModelPlugin):
+    name = "FastRandomizedLipschitz"
+    def __init__(self, proxies=10, dimensions=768, dist="euclidean", 
+                            rand_seed=12344, neighbors = 10, proxy_epochs=3, 
+                            perturbation_neighbors = 10, iterations = 25, batch_size = 256, 
+                            alpha = 0.5, feature_file = None, classifier_access = "classifier", **kwargs):
+        super().__init__(proxies = proxies, dimensions=dimensions, dist=dist, 
+                            rand_seed = rand_seed, neighbors = neighbors, proxy_epochs=proxy_epochs, 
+                            perturbation_neighbors=perturbation_neighbors, iterations = iterations, batch_size = batch_size, 
+                            alpha = alpha, feature_file = feature_file, classifier_access = classifier_access)
+
+
+    def build_plugin(self, **kwargs):
+        self.proxies = kwargs.get("proxies")
+        self.rand_seed = kwargs.get("rand_seed", 12344)
+        self.dimensions = kwargs.get("dimensions")
+        self.dist = kwargs.get("dist", "euclidean")    # TODO
+        self.proxy_epochs = kwargs.get("proxy_epochs", 3)
+        self.neighbors = kwargs.get("neighbors", 10)
+        self.perturbation_neighbors = kwargs.get("perturbation_neighbors", self.neighbors)
+
+        self.batch_size = kwargs.get("batch_size")
+        self.iterations = kwargs.get("iterations")
+        self.alpha = kwargs.get("alpha")
+
+        output_file = kwargs.get("feature_file")
+        if output_file is None:
+            import uuid
+            output_file = "-".join(["FRL", "features", uuid.uuid1().urn.split("-")[0].split(":")[-1]])
+        self._output_file = output_file + ".h5"
+        if os.path.exists(self._output_file):
+            os.remove(self._output_file)
+        self._writer = h5py.File(self._output_file, "w")  #we will delete old training features file
+        self._written = False
+        self._prev_idx = -1
+
+        self.epoch_count = 0
+        self.pre_epoch_flag = False
+        self.pre_epoch_num = -1
+        self.post_epoch_num = -1
+        self.activated = False
+        self.kdcluster = None
+        self.proxy_stage = True
+        self.lipschitz_stage = False
+
+        self.classifier_access = kwargs.get("classifier_access")
+        self._classifier_setup = False
+        self._classifier = None
+
+        
+        self._dist_func = self.build_dist(self.dist)
+        if self.dist == "euclidean":
+            self._preprocess = lambda x:x
+        elif self.dist == "cosine":
+            self._preprocess = torch.nn.functional.normalize
+
+        self._closest_features : List[SortedKeyList] = None
+
+
+
+    def build_dist(self, dist="euclidean"):
+        if dist == "euclidean":
+            return self.l2dist
+        elif dist == "cosine":
+            return self.cosdist
+        else:
+            raise ValueError("Invalid value for dist: %s"%dist)
+
+    def l2dist(self, x):
+        return torch.sqrt(((self.cluster_means - x)**2).sum(1)) # torch.argmin()
+    def cosdist(self, x):   #https://stackoverflow.com/questions/46409846/using-k-means-with-cosine-similarity-python
+        return self.l2dist(torch.nn.functional.normalize(x, dim=0))
+        
+    
+    def build_params(self, **kwargs):
+        self.cluster_means = torch.rand(self.proxies, self.dimensions)
+        self.inertia = 0
+        self.cluster_counts = torch.zeros(self.proxies)
+        self.kdcluster = None
+        self.high_density_thresholds = []
+
+        self.lipschitz_threshold = None
+        self.lipschitz_threshold_mean = None
+        self.smooth_lipschitz_threshold = None
+        self.smooth_lipschitz_threshold_mean = None
+        self.epsilon = None
+        self.proxy_label = []
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self._classifier_setup:
+            self._classifier = {mname:mlayer for mname, mlayer in model.named_modules()}[self.classifier_access]
+            self._classifier_setup = True
+        if not self.activated:
+            # perform the training here
+            if self.proxy_stage:
+                self.save_features(features)
+            if self.lipschitz_stage:
+                # this epoch, we are trying to find the m closest points to each cluster mean...
+                # for the batch, find each points distance to cluster means...??????
+                self.search_minibatch_kmeans_clusters(features)
+                
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            lscore, smoothlscore, proxy_labels = self.compute_lscores(features, feature_logits, model)
+            # here, we use the features, perturb them, compute the logits, then compute l_score, then compare, etc, etc
+            #dist, labels = self.compute_labels(features)
+            return feature_logits, features, secondary_outputs, kwargs, {"l_score": lscore, "smooth_l_score": smoothlscore, 
+                                                                            "l_threshold": [self.lipschitz_threshold]*features.shape[0], 
+                                                                            "smooth_l_threshold": [self.smooth_lipschitz_threshold]*features.shape[0],
+                                                                            "proxy_label": proxy_labels}
+
+
+    def compute_lscores(self, features, feature_logits, model):
+        perturbed = self.generate_perturbation(self.epsilon,self.dimensions,self.perturbation_neighbors).T.cuda()
+        lscore = [None]*features.shape[0]
+        smooth_lscore = [None]*features.shape[0]
+        for j,x in enumerate(features):
+            perturbation = x+perturbed
+            with torch.no_grad():
+                perturbed_logits = self._classifier(perturbation)
+                raw_logits = feature_logits[j].unsqueeze(0)
+
+                feature_lscore = torch.sqrt((perturbed**2).sum(1))
+                logit_lscore = torch.sqrt(((perturbed_logits - raw_logits)**2).sum(1))
+
+                smoothlogit_lscore = perturbed_logits[:,torch.argmax(raw_logits)] - raw_logits[:,torch.argmax(raw_logits)]
+
+                l_scores = logit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+                smooth_lscores = smoothlogit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+
+                lscore[j] = torch.max(l_scores).cpu().item()
+                smooth_lscore[j] = torch.max(smooth_lscores).cpu().item()
+        
+        _, idx = self.kdcluster.query(self._preprocess(features.cpu()), k=1, return_distance=True)   #.squeeze()
+        
+        return lscore, smooth_lscore, self.proxy_label[idx[:,0]]
+
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        if not self.activated:  # If already activated, we do not need to change anything
+            
+            if self.proxy_stage:    # Check if we are ready to move to lipschitz stage.
+                self.post_epoch_num = epoch    # Maybe have a better check in case something drops or epochs get skipped...
+                if self.pre_epoch_flag: # Means we had a pre_epoch flag so an epoch was completed.
+                    self.pre_epoch_flag = False # reset the flag
+                    self.epoch_count+=1
+                    if self.post_epoch_num != self.pre_epoch_num:
+                        raise ValueError("Epoch may have been skipped. Before: %i\tAfter: %i"%(self.pre_epoch_num, self.post_epoch_num))
+                self.lipschitz_stage = self.epoch_count >= self.proxy_epochs # better way to deal with this whole situation, i.e. once activated, never changes...!
+                self.proxy_stage = not self.lipschitz_stage
+                if self.proxy_stage:
+                    self._logger.info("RandomizedLipschitz continuing proxy stage")
+                else:
+
+                    try:    # in case it is closed already somehow...
+                        self._writer.close()
+                    except:
+                        pass
+                    
+                    self.performkmeans()
+                    self.highdensitybins(model)
+
+                    with torch.no_grad():
+                        self.proxy_label = torch.argmax(self._classifier(self.cluster_means.cuda()).cpu(), dim=1)
+                    self.high_density_thresholds  = torch.tensor(self.high_density_thresholds)
+                    
+                    self._logger.info("RandomizedLipschitz starting Lipschitz stage")
+            elif self.lipschitz_stage:
+                # check if we are done and can activate 
+                self._logger.info("RandomizedLipschitz has completed Lipschitz stage. Computing L values.")
+                lthresh = [0]*len(self.cluster_means)
+                lthreshmean = [0]*len(self.cluster_means)
+                smooththresh = [0]*len(self.cluster_means)
+                smooththreshmean = [0]*len(self.cluster_means)
+                with torch.no_grad():
+                    for idx in range(len(self.cluster_means)):
+                        raw_logits = self._classifier(self.cluster_means[idx].unsqueeze(0).cuda()).cpu()
+                        lipschitz_logits = self._classifier(torch.stack([item[0] for item in self._closest_features[idx]]).cuda()).cpu()
+
+                        # use euclidean here, only!!!!!!
+                        feature_lscore = torch.sqrt(((torch.stack([item[0] for item in self._closest_features[idx]])  - self.cluster_means[idx].unsqueeze(0))**2).sum(1))
+                        logit_lscore = torch.sqrt(((lipschitz_logits - raw_logits)**2).sum(1))
+
+                        smoothlogit_lscore = lipschitz_logits[:,torch.argmax(raw_logits)] - raw_logits[:,torch.argmax(raw_logits)]
+
+                        l_scores = logit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+                        smooth_lscores = smoothlogit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+
+                        lthresh[idx] = torch.max(l_scores).item()
+                        lthreshmean[idx] = torch.mean(l_scores).item()
+                        smooththresh[idx] = torch.max(smooth_lscores).item()
+                        smooththreshmean[idx] = torch.mean(smooth_lscores).item()
+                
+                self.lipschitz_threshold = max(lthresh)
+                self.lipschitz_threshold_mean = max(lthreshmean)
+                self.smooth_lipschitz_threshold = max(smooththresh)
+                self.smooth_lipschitz_threshold_mean = max(smooththreshmean)
+
+                self.epsilon = 1. / self.smooth_lipschitz_threshold
+
+                self.activated = True
+
+                # So we are done settng up neighbors...
+                # here, we have access to the model. We can not pass all our features and stuff through the model, get the logits, compute L, etc, etc
+                
+
+        if self.lipschitz_stage and self._closest_features is None and not self.activated:
+            """
+            if self.dist == "euclidean":
+                self.kdcluster = KDTree(self.cluster_means)
+            elif self.dist == "cosine":
+                self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+            else:
+                raise ValueError("Invalid value for dist: %s"%self.dist)
+            """
+            #model.classifier(self.cluster_means.cuda())
+            #self.labels = torch.argmax(model.classifier(self.cluster_means.cuda()), dim=1).cpu()   # potentially unneeded
+            # build the array of closest neighbor features...
+            self._closest_features = [SortedKeyList(key=lambda x:x[1]) for _ in range(len(self.cluster_means))]  #x <- (feature, distance)
+
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
+
+    def save_features(self, features):
+        feats = features.cpu().numpy()
+        if self._written:
+            self._writer["features"].resize((self._writer["features"].shape[0] + feats.shape[0]), axis=0)
+            self._writer["features"][-feats.shape[0]:] = feats        
+        else:   # First time writing -- we will need to create the dataset.
+            self._writer.create_dataset("features", data=feats, compression = "gzip", chunks=True, maxshape=(None,feats.shape[1]))
+            self._written = True
+
+    def performkmeans(self):
+        import time
+        inertia = []
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        
+        print("Starting KMeans for k={kval}, with {iters} iterations".format(kval=self.proxies, iters=self.iterations))
+        kmeans = MiniBatchKMeans(n_clusters = self.proxies, random_state = self.rand_seed, batch_size = self.batch_size)
+        stime = time.time()
+        for iters in range(self.iterations):
+            for i in range(0, data_size, self.batch_size):
+                current_data = data['features'][i:i+self.batch_size]
+                kmeans.partial_fit(current_data)
+            if iters%5 == 0:
+                etime = round(time.time() - stime, 2)
+                print("\t[{elapse} s] -- Completed {iters} iterations".format(iters=iters, elapse = etime))
+                stime = time.time()
+        print("\tCompeted MBKM for k={kval}, with inertia: {inertia}".format(kval=self.proxies, inertia = kmeans.inertia_))
+        data.close()
+        self.cluster_means = torch.tensor(kmeans.cluster_centers_)
+        self.inertia = kmeans.inertia_
+        return inertia
+
+    def highdensitybins(self, model: ModelAbstract):
+        if self.dist == "euclidean":
+            self.kdcluster = KDTree(self.cluster_means)
+        elif self.dist == "cosine":
+            self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+        else:
+            raise ValueError("Invalid value for dist: %s"%self.dist)
+
+        data = h5py.File(self._output_file, 'r')
+        data_size = data['features'].shape[0]
+        distance_bins = [[] for _ in range(self.proxies)]
+        print("Starting High Density estimation")
+        for i in range(0, data_size, self.batch_size):
+            current_data = data['features'][i:i+self.batch_size]
+            dist, indices = self.kdcluster.query(self._preprocess(torch.tensor(current_data)), k=1, return_distance=True)   #.squeeze()
+            for idx,val in enumerate(indices[:,0]):
+                distance_bins[val].append(dist[idx, 0])
+        
+        self.high_density_thresholds = [None]*self.proxies
+        import numpy as np
+        for proxy in range(self.proxies):
+            self.high_density_thresholds[proxy] = np.percentile(distance_bins[proxy], self.alpha * 100)
+        print("Completed High Density threshold estimation")
+        data.close()
+
+
+    def search_minibatch_kmeans_clusters(self, batch: torch.Tensor, ):
+        #batch_device = batch.get_device()   # TODO fix this hacky with the hack from https://stackoverflow.com/questions/58926054/how-to-get-the-device-type-of-a-pytorch-module-conveniently 2nd answer
+        local_batch = batch.cpu()
+        
+        for _, x in enumerate(local_batch):
+            resp = self._dist_func(x)   # distance of x to each cluster_center
+            # basically --> 
+            [feature_list.add((x, resp[idx].item())) for idx, feature_list in enumerate(self._closest_features)]
+        
+        # prune each to be length of self.neighbors...
+        self._closest_features = [SortedKeyList(feature_list[:self.neighbors], key=feature_list.key) for idx, feature_list in enumerate(self._closest_features)]
+        
+        # For each cluster center, we have a dict of closest features...
+
+    def generate_perturbation(self, epsilon, n_dims, n_samples):
+        Y = np.random.multivariate_normal(mean=[0], cov=np.eye(1,1), size=(n_dims, n_samples))
+        Y = np.squeeze(Y, -1)
+        Y /= np.sqrt(np.sum(Y * Y, axis=0))
+        U = np.random.uniform(low=0, high=1, size=(n_samples)) ** (1/n_dims)
+        Y *= U * epsilon # in my case radius is one
+        return torch.Tensor(Y)
\ No newline at end of file
diff --git a/src/ednaml/plugins/KMeansProxy.py b/src/ednaml/plugins/KMeansProxy.py
new file mode 100644
index 00000000..30797c53
--- /dev/null
+++ b/src/ednaml/plugins/KMeansProxy.py
@@ -0,0 +1,119 @@
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+
+class KMeansProxy(ModelPlugin):
+    name = "KMeansProxy"
+    def __init__(self, num_clusters=10, dimensions=768, dist="euclidean", rand_seed=145845, epochs=3, **kwargs):
+        super().__init__(num_clusters = num_clusters, dimensions=dimensions, dist=dist, rand_seed = rand_seed, epochs=epochs)
+
+
+    def build_plugin(self, **kwargs):
+        self.num_clusters = kwargs.get("num_clusters")
+        self.rand_seed = kwargs.get("rand_seed", 145845)
+        self.dimensions = kwargs.get("dimensions")
+        self.dist = kwargs.get("dist", "euclidean")    # TODO
+        self.epochs = kwargs.get("epochs", 3)
+        self.epoch_count = 0
+        self.pre_epoch_flag = False
+        self.pre_epoch_num = -1
+        self.post_epoch_num = -1
+        self.activated = False
+        self.kdcluster = None
+        self._dist_func = self.build_dist(self.dist)
+        if self.dist == "euclidean":
+            self._preprocess = lambda x:x
+        elif self.dist == "cosine":
+            self._preprocess = torch.nn.functional.normalize
+
+    def build_dist(self, dist="euclidean"):
+        if dist == "euclidean":
+            return self.l2dist
+        elif dist == "cosine":
+            return self.cosdist
+        else:
+            raise ValueError("Invalid value for dist: %s"%dist)
+
+    def l2dist(self, x):
+        return torch.argmin(torch.sqrt(((self.cluster_means - x)**2).sum(1)))
+    def cosdist(self, x):   #https://stackoverflow.com/questions/46409846/using-k-means-with-cosine-similarity-python
+        return self.l2dist(torch.nn.functional.normalize(x, dim=0))
+        
+    
+    def build_params(self, **kwargs):
+        self.cluster_means = torch.rand(self.num_clusters, self.dimensions)
+        self.cluster_counts = torch.zeros(self.num_clusters)
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self.activated:
+            # perform the training here
+            self.update_minibatch_kmeans_clusters(features)
+            
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            dist, labels = self.compute_labels(features)
+            return feature_logits, features, secondary_outputs, kwargs, {"cluster_mean": labels, "distance": dist}
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        if not self.activated:  # If already activated, we do not need to change anything
+            self.post_epoch_num = epoch    # Maybe have a better check in case something drops or epochs get skipped...
+            if self.pre_epoch_flag: # Means we had a pre_epoch flag so an epoch was completed.
+                self.pre_epoch_flag = False # reset the flag
+                self.epoch_count+=1
+                if self.post_epoch_num != self.pre_epoch_num:
+                    raise ValueError("Epoch may have been skipped. Before: %i\tAfter: %i"%(self.pre_epoch_num, self.post_epoch_num))
+            self.activated = self.epoch_count > self.epochs # better way to deal with this whole situation, i.e. once activated, never changes...!
+        if self.activated and self.kdcluster is None:
+            self._logger.info("KMP is active. Generating kdcluster")
+            if self.dist == "euclidean":
+                self.kdcluster = KDTree(self.cluster_means)
+            elif self.dist == "cosine":
+                self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+            else:
+                raise ValueError("Invalid value for dist: %s"%self.dist)
+
+            #model.classifier(self.cluster_means.cuda())
+            with torch.no_grad():
+                self.labels = torch.argmax(model.classifier(self.cluster_means.cuda()), dim=1).cpu()
+
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
+
+    def update_minibatch_kmeans_clusters(self, batch: torch.Tensor, ):
+        #batch_device = batch.get_device()   # TODO fix this hacky with the hack from https://stackoverflow.com/questions/58926054/how-to-get-the-device-type-of-a-pytorch-module-conveniently 2nd answer
+        local_batch = batch.cpu()
+        V = torch.zeros(self.num_clusters)
+        idxs = torch.empty(batch.shape[0], dtype=torch.int)
+        for j, x in enumerate(local_batch):
+            idxs[j] = self._dist_func(x)
+
+        # Update centers:
+        for j, x in enumerate(local_batch):
+            V[idxs[j]] += 1
+            eta = 1.0 / V[idxs[j]]
+            self.cluster_means[idxs[j]] = (1.0 - eta) * self.cluster_means[idxs[j]] + eta * x
+
+    def compute_labels(self, batch):
+        """Compute the cluster labels for dataset X given centers C.
+        """
+        # labels = np.argmin(pairwise_distances(C, X), axis=0) # THIS REQUIRES TOO MUCH MEMORY FOR LARGE X
+        q_batch = batch.cpu()
+        dist, idx = self.kdcluster.query(self._preprocess(q_batch), k=1, return_distance=True)   #.squeeze()
+        # TODO convert idx to the actual cluster means to the actual cluster labels...
+        return torch.from_numpy(dist).squeeze(1), torch.stack([self.labels[item[0]] for item in idx])
+
+
+
+"""cluster_distances = torch.zeros(self.num_clusters)
+    for cluster in range(self.num_clusters):
+        cluster_distances[cluster] = sum(torch.sqrt((features - self.cluster_means[cluster])**2))
+    c = torch.argmin(cluster_distances)
+    self.cluster_counts[c] += 1
+    self.cluster_means[c] += 1.0/self.cluster_counts[c]*(features - self.cluster_means[c])
+
+"""
\ No newline at end of file
diff --git a/src/ednaml/plugins/LogitConfidence.py b/src/ednaml/plugins/LogitConfidence.py
new file mode 100644
index 00000000..c6e88446
--- /dev/null
+++ b/src/ednaml/plugins/LogitConfidence.py
@@ -0,0 +1,62 @@
+import ednaml.core.decorators as edna
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+from sklearn.cluster import MiniBatchKMeans
+import h5py, os
+from ednaml.plugins.ModelPlugin import ModelPlugin
+
+class LogitConfidence(ModelPlugin):
+    name = "LogitConfidence"
+    def __init__(self, num_classes = 2, **kwargs):
+        super().__init__(num_classes = num_classes)
+
+    def build_plugin(self, **kwargs):
+        self.epochs = 1
+        self.activated = False
+        self.num_classes = kwargs.get("num_classes")
+
+    def build_params(self, **kwargs):
+        self.logit_confidence_logits = torch.zeros((1,self.num_classes))
+        self.logit_confidence_count = torch.zeros((1,self.num_classes))
+        self.logit_confidence = torch.zeros((1,self.num_classes))
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self.activated:
+            # perform the training here
+            self.add_to_average(feature_logits)
+            
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            logit, threshold = self.compute_labels(feature_logits)
+            return feature_logits, features, secondary_outputs, kwargs, {"logit": logit, "logit_threshold": threshold}
+
+    def add_to_average(self, feature_logits):
+        # Basically, for each class, find the entries where they are the max, average those
+        softlogits = torch.nn.functional.softmax(feature_logits.cpu(), dim=1)
+        one_hot = torch.nn.functional.one_hot(torch.argmax(softlogits, dim=1))
+        self.logit_confidence_logits += torch.sum(softlogits * one_hot, dim=0)
+        self.logit_confidence_count += torch.sum(one_hot, dim=0)
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        if not self.activated:  # If already activated, we do not need to change anything
+            # No need to do anything...
+            self.logit_confidence = self.logit_confidence_logits / self.logit_confidence_count
+            self.activated = True
+        self.pre_epoch_flag = False
+
+    
+    def compute_labels(self, feature_logits):
+        """Compute the cluster labels for dataset X given centers C.
+        """
+        # labels = np.argmin(pairwise_distances(C, X), axis=0) # THIS REQUIRES TOO MUCH MEMORY FOR LARGE X
+        soft_maxes = torch.max(torch.nn.functional.softmax(feature_logits.cpu(), dim=1),dim=1)
+        return soft_maxes[0], self.logit_confidence[0,soft_maxes[1]]
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
+
diff --git a/src/ednaml/plugins/ModelPlugin.py b/src/ednaml/plugins/ModelPlugin.py
new file mode 100644
index 00000000..b1690345
--- /dev/null
+++ b/src/ednaml/plugins/ModelPlugin.py
@@ -0,0 +1,68 @@
+import copy
+from logging import Logger
+from torch import nn
+class ModelPlugin():
+    name: str = "ModelPlugin"
+    _logger: Logger = None
+    activated: bool = False
+    def __init__(self, **kwargs):
+        #super().__init__()
+        self.build_plugin(**kwargs)
+        self.build_params(**kwargs)
+        
+        
+    def build_plugin(self, **kwargs):
+        """Sets up plugin-specific variables
+        """
+        pass
+
+    def build_params(self, **kwargs):
+        """Sets up plugin operational parameters
+        """
+        pass
+
+    def pre_forward(self, x, **kwargs):
+        """Called before the forward pass of a ModelAbstract, to modify the inputs if needed and to add to the secondary outputs
+
+        Args:
+            x (_type_): _description_
+        """
+        return x, kwargs, {}
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        """Called after the forward pass of a ModelAbstract, to use Plugin operations parameters to add to secondary outputs and adjust the outputs
+
+        Args:
+            x (_type_): _description_
+            feature_logits (_type_): _description_
+            features (_type_): _description_
+            secondary_outputs (_type_): _description_
+        """
+        return feature_logits, features, secondary_outputs, kwargs, {}
+
+    def pre_epoch(self, model: 'ModelAbstract', epoch: int = 0, **kwargs):
+        """Hook that is executed during the model training by BaseTrainer. Occurs before an epoch starts
+
+        Args:
+            epoch (int, optional): _description_. Defaults to 0.
+        """
+        pass
+
+    def post_epoch(self, model: 'ModelAbstract', epoch: int = 0, **kwargs):
+        """Hook that is executed during model training by BaseTrainer. Occurs at the end of an epoch.
+
+        Args:
+            epoch (int, optional): _description_. Defaults to 0.
+        """
+        pass
+
+    def save(self):
+        return copy.deepcopy({item:self.__dict__[item] for item in self.__dict__.keys() if item[0] != "_"})
+        #builtin_keys = [item for item in sdict.keys() if item[0] == "_"]
+        #[sdict.pop(key) for key in builtin_keys]
+        #return sdict
+        # returns an object of itself for saving
+
+    def load(self, save_dict_or_path):
+        # TODO adjust this so plugin can be loaded from a file, instead of directly passing parameters...
+        self.__dict__.update(save_dict_or_path) # probably need a better way to do this, i.e. iterate through dict, and keep what is ther, warn about extra keys, etc...
\ No newline at end of file
diff --git a/src/ednaml/plugins/RandomizedLipschitz.py b/src/ednaml/plugins/RandomizedLipschitz.py
new file mode 100644
index 00000000..872cf0c5
--- /dev/null
+++ b/src/ednaml/plugins/RandomizedLipschitz.py
@@ -0,0 +1,222 @@
+
+from typing import List
+from ednaml.models.ModelAbstract import ModelAbstract
+from ednaml.plugins import ModelPlugin
+import torch
+import torch.nn.functional
+from sklearn.neighbors import KDTree
+from ednaml.plugins.ModelPlugin import ModelPlugin
+from sortedcontainers import SortedKeyList
+import numpy as np
+
+class RandomizedLipschitz(ModelPlugin):
+    name = "RandomizedLipschitz"
+    def __init__(self, proxies=10, dimensions=768, dist="euclidean", rand_seed=12344, neighbors = 10, proxy_epochs=3, perturbation_neighbors = 10, **kwargs):
+        super().__init__(proxies = proxies, dimensions=dimensions, dist=dist, rand_seed = rand_seed, neighbors = neighbors, proxy_epochs=proxy_epochs, perturbation_neighbors=perturbation_neighbors)
+
+
+    def build_plugin(self, **kwargs):
+        self.proxies = kwargs.get("proxies")
+        self.rand_seed = kwargs.get("rand_seed", 12344)
+        self.dimensions = kwargs.get("dimensions")
+        self.dist = kwargs.get("dist", "euclidean")    # TODO
+        self.proxy_epochs = kwargs.get("proxy_epochs", 3)
+        self.neighbors = kwargs.get("neighbors", 10)
+        self.perturbation_neighbors = kwargs.get("perturbation_neighbors", self.neighbors)
+
+        self.epoch_count = 0
+        self.pre_epoch_flag = False
+        self.pre_epoch_num = -1
+        self.post_epoch_num = -1
+        self.activated = False
+        self.kdcluster = None
+        self.proxy_stage = True
+        self.lipschitz_stage = False
+
+        self.lipschitz_threshold = None
+        self.lipschitz_threshold_mean = None
+        self.smooth_lipschitz_threshold = None
+        self.smooth_lipschitz_threshold_mean = None
+        self.epsilon = None
+
+        
+        self._dist_func = self.build_dist(self.dist)
+        if self.dist == "euclidean":
+            self._preprocess = lambda x:x
+        elif self.dist == "cosine":
+            self._preprocess = torch.nn.functional.normalize
+
+        self._closest_features : List[SortedKeyList] = None
+
+
+
+    def build_dist(self, dist="euclidean"):
+        if dist == "euclidean":
+            return self.l2dist
+        elif dist == "cosine":
+            return self.cosdist
+        else:
+            raise ValueError("Invalid value for dist: %s"%dist)
+
+    def l2dist(self, x):
+        return torch.sqrt(((self.cluster_means - x)**2).sum(1)) # torch.argmin()
+    def cosdist(self, x):   #https://stackoverflow.com/questions/46409846/using-k-means-with-cosine-similarity-python
+        return self.l2dist(torch.nn.functional.normalize(x, dim=0))
+        
+    
+    def build_params(self, **kwargs):
+        self.cluster_means = torch.rand(self.proxies, self.dimensions)
+        self.cluster_counts = torch.zeros(self.proxies)
+
+
+    def post_forward(self, x, feature_logits, features, secondary_outputs, model, **kwargs):
+        if not self.activated:
+            # perform the training here
+            if self.proxy_stage:
+                self.update_minibatch_kmeans_clusters(features)
+            if self.lipschitz_stage:
+                # this epoch, we are trying to find the m closest points to each cluster mean...
+                # for the batch, find each points distance to cluster means...??????
+                self.search_minibatch_kmeans_clusters(features)
+                
+            return feature_logits, features, secondary_outputs, kwargs, {}
+        else:
+            lscore, smoothlscore = self.compute_lscores(features, feature_logits, model)
+            # here, we use the features, perturb them, compute the logits, then compute l_score, then compare, etc, etc
+            #dist, labels = self.compute_labels(features)
+            return feature_logits, features, secondary_outputs, kwargs, {"l_score": lscore, "smooth_l_score": smoothlscore, "l_threshold": [self.lipschitz_threshold]*features.shape[0], "smooth_l_threshold": [self.smooth_lipschitz_threshold]*features.shape[0]}
+
+
+    def compute_lscores(self, features, feature_logits, model):
+        perturbed = self.generate_perturbation(self.epsilon,self.dimensions,self.perturbation_neighbors).T.cuda()
+        lscore = [None]*features.shape[0]
+        smooth_lscore = [None]*features.shape[0]
+        for j,x in enumerate(features):
+            perturbation = x+perturbed
+            with torch.no_grad():
+                perturbed_logits = model.classifier(perturbation)
+                raw_logits = feature_logits[j].unsqueeze(0)
+
+                feature_lscore = torch.sqrt((perturbed**2).sum(1))
+                logit_lscore = torch.sqrt(((perturbed_logits - raw_logits)**2).sum(1))
+
+                smoothlogit_lscore = perturbed_logits[:,torch.argmax(raw_logits)] - raw_logits[:,torch.argmax(raw_logits)]
+
+                l_scores = logit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+                smooth_lscores = smoothlogit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+
+                lscore[j] = torch.max(l_scores).cpu().item()
+                smooth_lscore[j] = torch.max(smooth_lscores).cpu().item()
+        
+        return lscore, smooth_lscore
+
+
+    def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        if not self.activated:  # If already activated, we do not need to change anything
+            
+            if self.proxy_stage:    # Check if we are ready to move to lipschitz stage.
+                self.post_epoch_num = epoch    # Maybe have a better check in case something drops or epochs get skipped...
+                if self.pre_epoch_flag: # Means we had a pre_epoch flag so an epoch was completed.
+                    self.pre_epoch_flag = False # reset the flag
+                    self.epoch_count+=1
+                    if self.post_epoch_num != self.pre_epoch_num:
+                        raise ValueError("Epoch may have been skipped. Before: %i\tAfter: %i"%(self.pre_epoch_num, self.post_epoch_num))
+                self.lipschitz_stage = self.epoch_count >= self.proxy_epochs # better way to deal with this whole situation, i.e. once activated, never changes...!
+                self.proxy_stage = not self.lipschitz_stage
+                if self.proxy_stage:
+                    self._logger.info("RandomizedLipschitz continuing proxy stage")
+                else:
+                    self._logger.info("RandomizedLipschitz starting Lipschitz stage")
+            elif self.lipschitz_stage:
+                # check if we are done and can activate 
+                self._logger.info("RandomizedLipschitz has completed Lipschitz stage. Computing L values.")
+                lthresh = [0]*len(self.cluster_means)
+                lthreshmean = [0]*len(self.cluster_means)
+                smooththresh = [0]*len(self.cluster_means)
+                smooththreshmean = [0]*len(self.cluster_means)
+                with torch.no_grad():
+                    for idx in range(len(self.cluster_means)):
+                        raw_logits = model.classifier(self.cluster_means[idx].unsqueeze(0).cuda()).cpu()
+                        lipschitz_logits = model.classifier(torch.stack([item[0] for item in self._closest_features[idx]]).cuda()).cpu()
+
+                        # use euclidean here, only!!!!!!
+                        feature_lscore = torch.sqrt(((torch.stack([item[0] for item in self._closest_features[idx]])  - self.cluster_means[idx].unsqueeze(0))**2).sum(1))
+                        logit_lscore = torch.sqrt(((lipschitz_logits - raw_logits)**2).sum(1))
+
+                        smoothlogit_lscore = lipschitz_logits[:,torch.argmax(raw_logits)] - raw_logits[:,torch.argmax(raw_logits)]
+
+                        l_scores = logit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+                        smooth_lscores = smoothlogit_lscore[feature_lscore > 0] / feature_lscore[feature_lscore > 0]
+                        
+                        lthresh[idx] = torch.max(l_scores).item()
+                        lthreshmean[idx] = torch.mean(l_scores).item()
+                        smooththresh[idx] = torch.max(smooth_lscores).item()
+                        smooththreshmean[idx] = torch.mean(smooth_lscores).item()
+                
+                self.lipschitz_threshold = max(lthresh)
+                self.lipschitz_threshold_mean = max(lthreshmean)
+                self.smooth_lipschitz_threshold = max(smooththresh)
+                self.smooth_lipschitz_threshold_mean = max(smooththreshmean)
+
+                self.epsilon = 1. / self.smooth_lipschitz_threshold
+
+                self.activated = True
+
+                # So we are done settng up neighbors...
+                # here, we have access to the model. We can not pass all our features and stuff through the model, get the logits, compute L, etc, etc
+                
+
+        if self.lipschitz_stage and self._closest_features is None and not self.activated:
+            """
+            if self.dist == "euclidean":
+                self.kdcluster = KDTree(self.cluster_means)
+            elif self.dist == "cosine":
+                self.kdcluster = KDTree(torch.nn.functional.normalize(self.cluster_means))
+            else:
+                raise ValueError("Invalid value for dist: %s"%self.dist)
+            """
+            #model.classifier(self.cluster_means.cuda())
+            #self.labels = torch.argmax(model.classifier(self.cluster_means.cuda()), dim=1).cpu()   # potentially unneeded
+            # build the array of closest neighbor features...
+            self._closest_features = [SortedKeyList(key=lambda x:x[1]) for _ in range(len(self.cluster_means))]  #x <- (feature, distance)
+
+
+    def pre_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
+        self.pre_epoch_flag = True
+        self.pre_epoch_num = epoch
+
+    def update_minibatch_kmeans_clusters(self, batch: torch.Tensor, ):
+        #batch_device = batch.get_device()   # TODO fix this hacky with the hack from https://stackoverflow.com/questions/58926054/how-to-get-the-device-type-of-a-pytorch-module-conveniently 2nd answer
+        local_batch = batch.cpu()
+        V = torch.zeros(self.proxies)
+        idxs = torch.empty(batch.shape[0], dtype=torch.int)
+        for j, x in enumerate(local_batch):
+            idxs[j] = torch.min(self._dist_func(x),0)[1].item()
+
+        # Update centers:
+        for j, x in enumerate(local_batch):
+            V[idxs[j]] += 1
+            eta = 1.0 / V[idxs[j]]
+            self.cluster_means[idxs[j]] = (1.0 - eta) * self.cluster_means[idxs[j]] + eta * x
+
+    def search_minibatch_kmeans_clusters(self, batch: torch.Tensor, ):
+        #batch_device = batch.get_device()   # TODO fix this hacky with the hack from https://stackoverflow.com/questions/58926054/how-to-get-the-device-type-of-a-pytorch-module-conveniently 2nd answer
+        local_batch = batch.cpu()
+        
+        for _, x in enumerate(local_batch):
+            resp = self._dist_func(x)   # distance of x to each cluster_center
+            # basically --> 
+            [feature_list.add((x, resp[idx].item())) for idx, feature_list in enumerate(self._closest_features)]
+        
+        # prune each to be length of self.neighbors...
+        self._closest_features = [SortedKeyList(feature_list[:self.neighbors], key=feature_list.key) for idx, feature_list in enumerate(self._closest_features)]
+        
+        # For each cluster center, we have a dict of closest features...
+
+    def generate_perturbation(self, epsilon, n_dims, n_samples):
+        Y = np.random.multivariate_normal(mean=[0], cov=np.eye(1,1), size=(n_dims, n_samples))
+        Y = np.squeeze(Y, -1)
+        Y /= np.sqrt(np.sum(Y * Y, axis=0))
+        U = np.random.uniform(low=0, high=1, size=(n_samples)) ** (1/n_dims)
+        Y *= U * epsilon # in my case radius is one
+        return torch.Tensor(Y)
\ No newline at end of file
diff --git a/src/ednaml/plugins/__init__.py b/src/ednaml/plugins/__init__.py
new file mode 100644
index 00000000..da297e5d
--- /dev/null
+++ b/src/ednaml/plugins/__init__.py
@@ -0,0 +1,1 @@
+from ednaml.plugins.ModelPlugin import ModelPlugin
\ No newline at end of file
diff --git a/src/ednaml/storage/AzureStorage.py b/src/ednaml/storage/AzureStorage.py
new file mode 100644
index 00000000..19029d88
--- /dev/null
+++ b/src/ednaml/storage/AzureStorage.py
@@ -0,0 +1,60 @@
+from typing import Dict
+try:
+    from azure.storage.blob import BlockBlobService
+    from azure.storage.blob import AppendBlobService
+except:
+    pass
+from ednaml.storage.BaseStorage import BaseStorage
+import os,shutil,gzip
+
+class AzureStorage(BaseStorage):        
+    def build_params(self, **kwargs):
+        self.storage_account_key = kwargs.get("storage_account_key")
+        self.storage_account = kwargs.get("storage_account", "")
+        self.containername = kwargs.get("containername", "default")
+        self.blobname = kwargs.get("blobname", "default")
+        self.local_azfilename = kwargs.get("local_filename", self.blobname) # "../Data/cars_datasets_models2"
+        self.uploadblobname = kwargs.get("uploadblobname","default")    # TODO this needs to be dealt with w.r.t. config saves...
+
+    def read(self):
+        blob_service_client_instance = BlockBlobService(account_name=self.storage_account, 
+            account_key=self.storage_account_key)
+
+        blob_service_client_instance.create_container(self.containername)
+        downloadpath = os.path.join("./",self.blobname)
+
+        print("\nDownloading blob to " + downloadpath)
+        blob_service_client_instance.get_blob_to_path( self.containername, self.blobname, downloadpath)
+        if not os.path.exists(self.local_azfilename): 
+            with gzip.open(downloadpath, 'rb') as f_in: 
+                with open(self.local_azfilename, 'wb') as f_out: 
+                    shutil.copyfileobj(f_in, f_out) 
+
+    def write(self, data):  # TODO should upload name be specified here as well...????
+        #print("Storage args = ", self.STORAGE_ARGS)
+        blob_service_client_instance = BlockBlobService(account_name=self.storage_account, 
+            credential=self.storage_account_key)
+
+        blob_client_instance_upload = blob_service_client_instance.get_blob_client(
+            self.containername, self.uploadblobname, snapshot=None)
+        
+        blob_client_instance_upload.upload_blob(data)   # TODO Possibly save data from somewhere. Is it a stream...need to read up on this
+
+        #blob_client_instance_upload.append_blob_from_text()
+
+    def append(self,data):
+        print("This is upload code in Azure")
+        append_blob_service = AppendBlobService(
+            account_name=self.storage_account, 
+            account_key=self.storage_account_key)
+        # Creates an append blob for this app.
+        # What does this do...?
+        append_blob_service.create_container( self.containername)
+        append_blob_service.create_blob( self.containername, 'log.txt')
+        append_blob_service.append_blob_from_text( self.containername, 'log.txt', data)
+
+    def copy(self,src):
+        print("Uploading/copying logs to cloud storage")    
+        with open(src, 'r') as content_file:
+            content = content_file.read()
+            self.append(content)
diff --git a/src/ednaml/storage/BaseStorage.py b/src/ednaml/storage/BaseStorage.py
new file mode 100644
index 00000000..8be5601b
--- /dev/null
+++ b/src/ednaml/storage/BaseStorage.py
@@ -0,0 +1,22 @@
+class BaseStorage:
+    storage_type: str
+    storage_url: str
+    def __init__(self, type, url, **kwargs):
+        self.storage_type = type
+        self.storage_url = url
+        self.build_params(**kwargs)
+        
+    def build_params(self, **kwargs):
+        pass
+
+    def read(self):
+        print("Base read call")
+
+    def write(self, data):
+        print("Base write call",data)
+
+    def append(self,data):
+        print("Append call",data)
+
+    def copy(self,src):
+        print("Copy call ",src)
diff --git a/src/ednaml/storage/LocalStorage.py b/src/ednaml/storage/LocalStorage.py
new file mode 100644
index 00000000..626315fd
--- /dev/null
+++ b/src/ednaml/storage/LocalStorage.py
@@ -0,0 +1,27 @@
+from ednaml.config.EdnaMLConfig import EdnaMLConfig
+from typing import Dict
+from azure.storage.blob import BlockBlobService
+from azure.storage.blob import AppendBlobService
+from ednaml.config.SaveConfig import SaveConfig
+from ednaml.storage.BaseStorage import BaseStorage
+import os,shutil,gzip
+
+class LocalStorage(BaseStorage):
+    def read(self):
+        print("Implement download file here")
+        downloadpath = "This is dummy name"
+        az_jsonfile = "../Data/cars_datasets_models2"
+        if not os.path.exists(az_jsonfile): 
+            with gzip.open(downloadpath, 'rb') as f_in: 
+                with open(az_jsonfile, 'wb') as f_out: 
+                    shutil.copyfileobj(f_in, f_out) 
+
+    def write(self, data):
+        print("Implement write file here")
+
+        #blob_client_instance_upload.append_blob_from_text()
+
+    def append(self,data):
+        print("Implement append file here")
+    def copy(self,src,dst):
+        print("Implement copy to storage here")
diff --git a/src/ednaml/storage/__init__.py b/src/ednaml/storage/__init__.py
new file mode 100644
index 00000000..fed8894c
--- /dev/null
+++ b/src/ednaml/storage/__init__.py
@@ -0,0 +1,2 @@
+from ednaml.storage.BaseStorage import BaseStorage
+from ednaml.storage.AzureStorage import AzureStorage
diff --git a/src/ednaml/trainer/BaseTrainer.py b/src/ednaml/trainer/BaseTrainer.py
index b148251b..e56c6660 100644
--- a/src/ednaml/trainer/BaseTrainer.py
+++ b/src/ednaml/trainer/BaseTrainer.py
@@ -1,6 +1,7 @@
 import json, logging, os, shutil
 from logging import Logger
-from typing import Dict, List
+from typing import Dict, List, Tuple
+from ednaml.config.StorageConfig import StorageConfig
 
 import torch
 from torch.utils.data import DataLoader
@@ -9,7 +10,7 @@
 from ednaml.crawlers import Crawler
 from ednaml.models.ModelAbstract import ModelAbstract
 from ednaml.utils.LabelMetadata import LabelMetadata
-
+from ednaml.storage.BaseStorage import BaseStorage
 
 class BaseTrainer:
     model: ModelAbstract
@@ -34,6 +35,7 @@ class BaseTrainer:
     metadata: Dict[str, str]
     labelMetadata: LabelMetadata
     logger: logging.Logger
+    storage: BaseStorage
 
     def __init__(
         self,
@@ -51,9 +53,9 @@ def __init__(
         crawler: Crawler,
         config: EdnaMLConfig,
         labels: LabelMetadata,
+        storage: BaseStorage,
         **kwargs
     ):
-
         self.model = model
         self.parameter_groups = list(self.model.parameter_groups.keys())
         self.loss_fn_order = {
@@ -102,6 +104,9 @@ def __init__(
         self.labelMetadata = labels
         self.crawler = crawler
         self.config = config
+        self.storage = storage
+        # Add later -- download data from Azure/other file instance
+        #self.downloadData()
 
         self.buildMetadata(
             # TODO This is not gonna work with the torchvision wrapper -- ned to fix that; because crawler is not set up for that pattern...?
@@ -113,12 +118,21 @@ def __init__(
         self.accumulation_count = 0
         self.evaluateFlag = False
         self.saveFlag = False
+        self.init_setup(**kwargs)
+
+    def init_setup(self, **kwargs):
+        pass
+
+    def downloadData(self): #TODO
+        #storageinstance = AzureStorage(self.cfg.STORAGE)
+        self.storage.read()
+
 
     def buildMetadata(self, **kwargs):
         for keys in kwargs:
             self.metadata[keys] = kwargs.get(keys)
 
-    def setup(
+    def apply(
         self,
         step_verbose: int = 5,
         save_frequency: int = 5,
@@ -148,9 +162,14 @@ def setup(
         self.gpus = gpus
 
         if self.gpus != 1:
-            raise NotImplementedError()
+            self.logger.warning("Multi-gpu or non-gpu not yet fully supported.")
+
 
-        self.model.cuda() # moves the model into GPU
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        if self.gpus:
+            #self.model.cuda() # moves the model into GPU
+            self.logger.info("%i GPUs available"%self.gpus)
+        self.model.to(self.device)
 
         self.fp16 = fp16
         # if self.fp16 and self.apex is not None:
@@ -197,6 +216,9 @@ def save(self, save_epoch: int = None):
             )
             for lossname in self.loss_scheduler
         }
+        # TODO need to check if ModelAbstract contains the plugins that came with it...
+        # if not, we will need to save plugins in a plugins artifact.
+        # And even if they did, we need to split up model and plugins, since they are not strictly part of the model.
 
         # save_dict["loss_optimizer"] = [self.loss_optimizer[idx].state_dict() if self.loss_optimizer[idx] is not None else None for idx in range(self.num_losses)]
         # save_dict["loss_scheduler"] = [self.loss_scheduler[idx].state_dict() if self.loss_scheduler[idx] is not None else None for idx in range(self.num_losses)]
@@ -227,12 +249,17 @@ def save(self, save_epoch: int = None):
         
         if self.config.SAVE.LOG_BACKUP:
             LOGGER_SAVE = os.path.join(self.backup_directory, self.logger_file)
+            #appending the logs to azure
+            #self.storage.copy(os.path.join(self.save_directory, self.logger_file)) # TODO
+            # This is the local path where we are copying data . This can be commented/removed later
             if os.path.exists(LOGGER_SAVE):
                 os.remove(LOGGER_SAVE)
             shutil.copy2(
                 os.path.join(self.save_directory, self.logger_file), LOGGER_SAVE
             )
 
+                
+
     def load(self, load_epoch):
         self.logger.info(
             "Resuming training from epoch %i. Loading saved state from %i"
@@ -339,14 +366,16 @@ def train(self, continue_epoch=0):
             self.initial_evaluate()
         else:
             self.logger.info("Skipping initial evaluation.")
-
+        
         self.logger.info("Starting training from %i" % continue_epoch)
         self.zeroGrad()
         self.evaluateFlag = False
         self.saveFlag = False
         for epoch in range(self.epochs + 1):
             if epoch >= continue_epoch:
+                # TODO pre epoch
                 self.epoch_step(epoch)
+                # TODO post epoch
             else:
                 self.global_epoch = epoch + 1
 
@@ -370,7 +399,7 @@ def epoch_step(self, epoch):
                 self.printOptimizerLearningRates()
 
             self.model.train() #train == we are tracking all numbers and computation graph
-
+            batch = self.move_to_device(batch)
             loss: torch.Tensor = self.step(batch) #perform function and returns loss
             loss = loss / self.accumulation_steps
             loss.backward()
@@ -422,10 +451,14 @@ def epoch_step(self, epoch):
             else:
                 self.save()
         self.global_epoch += 1
+    
+    def move_to_device(self, batch) -> Tuple[torch.Tensor]:
+        return (item.to(self.device) for item in batch)
 
     def step(self, batch) -> torch.Tensor:
         # compute the loss, and return it
         # print("!!!!!!!!!! batch !!!!!!!!!!!!!!!",batch)
+
         raise NotImplementedError()
 
     def updateGradients(self):
diff --git a/src/ednaml/trainer/ClassificationTrainer.py b/src/ednaml/trainer/ClassificationTrainer.py
index ae4e1863..675481db 100644
--- a/src/ednaml/trainer/ClassificationTrainer.py
+++ b/src/ednaml/trainer/ClassificationTrainer.py
@@ -4,61 +4,20 @@
 import os
 import torch
 import numpy as np
-import ednaml.loss.builders
-from typing import List
-from ednaml.crawlers import Crawler
 from ednaml.trainer import BaseTrainer
-from ednaml.utils.LabelMetadata import LabelMetadata
 
 
 class ClassificationTrainer(BaseTrainer):
-    def __init__(
-        self,
-        model: torch.nn.Module,
-        loss_fn: List[ednaml.loss.builders.LossBuilder],
-        optimizer: torch.optim.Optimizer,
-        loss_optimizer: List[torch.optim.Optimizer],
-        scheduler: torch.optim.lr_scheduler._LRScheduler,
-        loss_scheduler: List[torch.optim.lr_scheduler._LRScheduler],
-        train_loader,
-        test_loader,
-        epochs: int,
-        skipeval,
-        logger,
-        crawler: Crawler,
-        config,
-        labels: LabelMetadata,
-        **kwargs
-    ):
-
-        super().__init__(
-            model,
-            loss_fn,
-            optimizer,
-            loss_optimizer,
-            scheduler,
-            loss_scheduler,
-            train_loader,
-            test_loader,
-            epochs,
-            skipeval,
-            logger,
-            crawler,
-            config,
-            labels,
-            **kwargs
-        )
-
+    def init_setup(self, **kwargs):
         self.softaccuracy = []
 
     # Steps through a batch of data
-    def step(self, batch): # Sanjyot --- simplify this
+    def step(self, batch): 
         batch_kwargs = {}
         (
             img,
             batch_kwargs["labels"],
-        ) = batch  # This is the tensor response from collate_fn
-        img, batch_kwargs["labels"] = img.cuda(), batch_kwargs["labels"].cuda()
+        ) = batch  
         # logits, features, labels
         batch_kwargs["logits"], batch_kwargs["features"], _ = self.model(img)
         batch_kwargs["epoch"] = self.global_epoch  # For CompactContrastiveLoss
@@ -97,7 +56,7 @@ def evaluate_impl(self):
                 self.test_loader, total=len(self.test_loader), leave=False
             ):
                 data, label = batch
-                data = data.cuda()
+                data = data.to(self.device)
                 logit, feature, _ = self.model(data)
                 feature = feature.detach().cpu()
                 logit = logit.detach().cpu()
@@ -125,7 +84,7 @@ def evaluate_impl(self):
         self.logger.info("Accuracy: {:.3%}".format(accuracy))
         self.logger.info("Micro F-score: {:.3f}".format(micro_fscore))
         self.logger.info("Weighted F-score: {:.3f}".format(weighted_fscore))
-        return logit_labels, labels, features
+        return logit_labels, labels, logits
 
     def saveMetadata(
         self,
diff --git a/src/ednaml/trainer/HFTrainer.py b/src/ednaml/trainer/HFTrainer.py
new file mode 100644
index 00000000..809dc633
--- /dev/null
+++ b/src/ednaml/trainer/HFTrainer.py
@@ -0,0 +1,118 @@
+import tqdm
+from sklearn.metrics import f1_score
+import torch
+import numpy as np
+from ednaml.trainer import BaseTrainer
+
+
+class HFTrainer(BaseTrainer):
+
+    def init_setup(self, **kwargs):
+        #self.maskedaccuracy = []
+        self.softaccuracy = []
+
+    def step(self, batch):
+        #batch = tuple(item.cuda() for item in batch)
+        (
+            all_input_ids,
+            all_attention_mask,
+            all_token_type_ids,
+            all_masklm,
+            all_annotations,
+            all_labels
+        ) = batch
+        outputs = self.model(
+            all_input_ids,
+            token_type_ids=all_token_type_ids,
+            attention_mask=all_attention_mask,
+            output_attentions = True,
+            output_hidden_states = True,
+            secondary_inputs=all_annotations       # NOT for HFTrainer! because it only expects specific inputs!
+        )
+
+        logits = outputs[0]
+        features = outputs[1]
+        secondaries = outputs[2]
+        
+
+        logits_loss = self.loss_fn["classification"](
+            logits=logits, labels=all_labels[:,0]   # Because there could be multiple labels...
+        )
+        softmax_accuracy = (
+                    (logits.max(1)[1] == all_labels[:,0])
+                    .float()
+                    .mean()
+                )
+
+        self.losses["classification"].append(logits_loss.item())
+        self.softaccuracy.append(softmax_accuracy.cpu().item())
+        
+        return logits_loss
+
+
+    def evaluate_impl(self):
+        logits, labels = [],[]
+        with torch.no_grad():
+            for batch in tqdm.tqdm(
+                self.test_loader, total=len(self.test_loader), leave=False
+            ):
+                batch = tuple(item.to(self.device) for item in batch)
+                (
+                    all_input_ids,
+                    all_attention_mask,
+                    all_token_type_ids,
+                    all_masklm,
+                    all_annotations,
+                    all_labels
+                ) = batch
+                outputs = self.model(
+                            all_input_ids,
+                            token_type_ids=all_token_type_ids,
+                            attention_mask=all_attention_mask,
+                            output_attentions = True,
+                            output_hidden_states = True,
+                            secondary_inputs=all_annotations
+                        )
+                raw_logits = outputs[0]
+                raw_features = outputs[1]
+                raw_secondaries = outputs[2]
+                logit = raw_logits.detach().cpu()
+                label = all_labels.detach().cpu()
+
+                logits.append(logit)
+                labels.append(label)
+
+
+        logits, labels = (
+            torch.cat(logits, dim=0),
+            torch.cat(labels, dim=0)    )
+        self.logger.info("Obtained logits and labels, validation in progress")
+
+
+        logit_labels = torch.argmax(logits, dim=1)
+        accuracy = (logit_labels == labels[:,0]).sum().float() / float(labels[:,0].size(0))
+        micro_fscore = np.mean(f1_score(labels[:,0], logit_labels, average="micro"))
+        weighted_fscore = np.mean(f1_score(labels[:,0], logit_labels, average="weighted"))
+        self.logger.info("\tAccuracy: {:.3%}".format(accuracy))
+        self.logger.info("\tMicro F-score: {:.3f}".format(micro_fscore))
+        self.logger.info("\tWeighted F-score: {:.3f}".format(weighted_fscore))
+
+        return logit_labels, labels, logits
+    
+    def printStepInformation(self):
+        loss_avg = 0.0
+        for lossname in self.losses:
+            loss_avg += (
+                sum(self.losses[lossname][-self.step_verbose :])
+                / self.step_verbose
+            )
+        #loss_avg /= self.num_losses
+        soft_avg = sum(self.softaccuracy[-100:]) / float(
+            len(self.softaccuracy[-100:])
+        )
+
+        self.logger.info(
+            "Epoch{0}.{1}\tClassification Loss: {2:.3f}\t Training Accuracy:  {3:.3f}".format(
+                self.global_epoch, self.global_batch, loss_avg, soft_avg
+                )
+        )
\ No newline at end of file
diff --git a/src/ednaml/trainer/MultiBranchTrainer.py b/src/ednaml/trainer/MultiBranchTrainer.py
index dc31ab5c..cb1b5c74 100644
--- a/src/ednaml/trainer/MultiBranchTrainer.py
+++ b/src/ednaml/trainer/MultiBranchTrainer.py
@@ -102,9 +102,9 @@ def step(self, batch):
             # Check if this is a loss for a soft target
             if lossname in self.model.soft_names:
                 # we need to adjust the labels...otherwise, we do not need to adjust labels...
-                akwargs["labels"] = batch_kwargs["logits"][
+                akwargs["labels"] = torch.argmax(batch_kwargs["logits"][
                     self.model_nameorder[self.model.soft_target_output_source]
-                ]
+                ], dim=1)
             else:
                 akwargs["labels"] = batch_kwargs["labels"][
                     :, self.data_labelorder[self.model_name_label_map[lossname]]
diff --git a/src/ednaml/trainer/__init__.py b/src/ednaml/trainer/__init__.py
index f8bd19a5..7a5d8631 100644
--- a/src/ednaml/trainer/__init__.py
+++ b/src/ednaml/trainer/__init__.py
@@ -3,3 +3,4 @@
 from ednaml.trainer.ClassificationTrainer import ClassificationTrainer
 from ednaml.trainer.MultiClassificationTrainer import MultiClassificationTrainer
 from ednaml.trainer.MultiBranchTrainer import MultiBranchTrainer
+from ednaml.trainer.HFTrainer import HFTrainer
diff --git a/src/ednaml/utils/__init__.py b/src/ednaml/utils/__init__.py
index 6bdbaefd..c5ad8d60 100644
--- a/src/ednaml/utils/__init__.py
+++ b/src/ednaml/utils/__init__.py
@@ -3,20 +3,23 @@
 import logging
 import sys
 from ednaml.exceptions import ErrorDuringImport
-
+import importlib.util
 
 def locate_class(
     package="ednaml",
     subpackage="core",
-    classpackage="EdnaML",
+    classpackage=None,
     classfile=None,
     forceload=0,
 ):
     """Locate an object by name or dotted path, importing as necessary."""
-    if classfile is None:
-        parts = package.split(".") + [subpackage, classpackage]
+    if classpackage is None:
+        parts = package.split(".") + [subpackage]
     else:
-        parts = package.split(".") + [subpackage, classfile, classpackage]
+        if classfile is None:
+            parts = package.split(".") + [subpackage, classpackage]
+        else:
+            parts = package.split(".") + [subpackage, classfile, classpackage]
     module, n = None, 0
     while n < len(parts):
         nextmodule = safeimport(".".join(parts[: n + 1]), forceload)
@@ -83,6 +86,14 @@ def safeimport(path, forceload=0, cache={}):
             return None
     return module
 
+def path_import(absolute_path):
+   '''implementation taken from https://docs.python.org/3/library/importlib.html#importing-a-source-file-directly'''
+   """And from https://stackoverflow.com/questions/41861427/python-3-5-how-to-dynamically-import-a-module-given-the-full-file-path-in-the"""
+   spec = importlib.util.spec_from_file_location(absolute_path, absolute_path)
+   module = importlib.util.module_from_spec(spec)
+   spec.loader.exec_module(module)
+   return module
+
 
 def extend_mean_arguments(
     params_to_fix: List[str] = [0.5, 0.5], channels=3
diff --git a/src/ednaml/utils/tokenizers/HFAutoTokenizer.py b/src/ednaml/utils/tokenizers/HFAutoTokenizer.py
new file mode 100644
index 00000000..160eecb6
--- /dev/null
+++ b/src/ednaml/utils/tokenizers/HFAutoTokenizer.py
@@ -0,0 +1,14 @@
+import warnings
+try:
+    from transformers import AutoTokenizer
+except:
+    AutoTokenizer = object
+    warnings.warn("`transformers` not installed. If using AutoTokenizer, there may be issues.")
+
+
+class HFAutoTokenizer:
+    def __init__(self, **kwargs):
+        raise RuntimeError("`__init__` called for HFAutoTokenizer. Use `__new__` instead")
+
+    def __new__(cls, from_pretrained: str, **kwargs):
+        return AutoTokenizer.from_pretrained(from_pretrained, **kwargs)
\ No newline at end of file
diff --git a/src/ednaml/utils/tokenizers/__init__.py b/src/ednaml/utils/tokenizers/__init__.py
index 513dccdd..ccfcf144 100644
--- a/src/ednaml/utils/tokenizers/__init__.py
+++ b/src/ednaml/utils/tokenizers/__init__.py
@@ -139,3 +139,5 @@ def _is_punctuation(char):
     AlbertBasicTokenizer,
     AlbertFullTokenizer,
 )
+
+from .HFAutoTokenizer import HFAutoTokenizer
diff --git a/stack/mongod.conf b/stack/mongo-db/mongod.conf
similarity index 94%
rename from stack/mongod.conf
rename to stack/mongo-db/mongod.conf
index c678bbff..97d9160b 100644
--- a/stack/mongod.conf
+++ b/stack/mongo-db/mongod.conf
@@ -28,7 +28,8 @@ net:
 processManagement:
   timeZoneInfo: /usr/share/zoneinfo
 
-#security:
+security:
+  authorization: enabled
 
 #operationProfiling:
 
diff --git a/stack/readme.md b/stack/readme.md
index 070983eb..d4da13ff 100644
--- a/stack/readme.md
+++ b/stack/readme.md
@@ -1,5 +1,12 @@
 # The EDNA stack
 
+
 MongoDB
     - config
-    
\ No newline at end of file
+    
+
+# What gets saved
+
+## EdnaML
+
+EdnaML during training
\ No newline at end of file
diff --git a/usage-docs/config-full.yml b/usage-docs/config-full.yml
index 0fa20e64..5a24fbdb 100644
--- a/usage-docs/config-full.yml
+++ b/usage-docs/config-full.yml
@@ -1,89 +1,197 @@
+# EXECUTION manages ML training and evaluation. Use with EdnaML
 EXECUTION:
+  # A trainer to use, from ednaml.trainers. A custom trainer can be implicitly added
+  TRAINER: BaseTrainer
+  # Arguments for the trainer
+  TRAINER_ARGS: 
+    accumulation_steps: 1
+  # An optimizer builder from ednaml.optimizers
   OPTIMIZER_BUILDER: ClassificationOptimizer
-  MODEL_SERVING: None
+  # Unused for now
+  MODEL_SERVING: Unused
+  # Number of epochs to train with 
   EPOCHS: 10
-  SKIPEVAL: True
-  TEST_FREQUENCY: 5
+  # Whether to skip the initial evaluation
+  SKIPEVAL: False
+  # The frequency, in epochs, for model testing
+  TEST_FREQUENCY: 1
+  # Unused
   FP16: False
+  # DATAREADER is the data-ingest unit that determines the data source. 
+  # It is defined in DATAREADER.DATAREADER below. A DataReader contains a 
+  # Crawler that yields raw samples, and a Generator that yields batches
+  # to a classifier. A generator may contain an associated 
+  # Dataset to manage sample preprocessing.
   DATAREADER: 
-    DATAREADER: VehicleColor
+    # A DataReader class from ednaml.datareaders for built-in data sources, e.g. TorchvisionDatareader for datasets from TorchVision
+    DATAREADER: DataReader
+    # Arguments for the DataReader's crawlers. The Crawler itself may be a custom Crawler, defined in a separate file. Custom Crawlers are implicitly added.
     CRAWLER_ARGS: {}
+    # Arguments for the DataReader's Dataset. The Dataset itself may be a custom Dataset, defined in a separate file. Custom Dataset are implicitly added.
     DATASET_ARGS: 
-      pathidx: 0
-      annotationidx: 1
-      classificationclass: color
+    # Whether to switch out the Generator in the DataReader with another built-in Generator. Note: custom Generators are implicitly added
     GENERATOR: 
+    # Arguments for the DataReader's Generator. The Generator itself may be a custom Generator, defined in a separate file. Custom Generator are implicitly added.
     GENERATOR_ARGS: 
-    DATALOADER_ARGS: {}
-  TRAINER: ClassificationTrainer
-  TRAINER_ARGS: 
-    accumulation_steps: 4
+  # PLUGIN manages how plugins are used. Unused for EXECUTION
+  PLUGIN:
+    # Controls when plugins are used. `always` means plugins are always used. `warmup` means plugins are used only until they are ready (i.e. activated)
+    # `activated` means only plugins that are ready are used.
+    HOOKS: always # always | warmup | activated
+    # Whether to reset already generated PLUGINS. Unused.
+    RESET: False
 
+# DEPLOYMENT manages ML model deployment or prediction. Use with EdnaDeploy
+DEPLOYMENT:
+  # A Deploy to use, from ednaml.deploy
+  DEPLOY: BaseDeploy
+  # Arguments for any output details for the Deploy. Refer to Deploy's documentation for specifics
+  OUTPUT_ARGS:  {}
+  # Arguments for the Deploy itself. Refer to Deploy's documentation for specifics.
+  DEPLOYMENT_ARGS: {}
+  # The number of epochs to run this Deployment.
+  EPOCHS: 1
+  # DATAREADER is the data-ingest unit that determines the data source. 
+  # It is defined in DATAREADER.DATAREADER below. A DataReader contains a 
+  # Crawler that yields raw samples, and a Generator that yields batches
+  # to a classifier. A generator may contain an associated 
+  # Dataset to manage sample preprocessing.
+  DATAREADER: 
+    # A DataReader class from ednaml.datareaders for built-in data sources, e.g. TorchvisionDatareader for datasets from TorchVision
+    DATAREADER: DataReader
+    # Arguments for the DataReader's crawlers. The Crawler itself may be a custom Crawler, defined in a separate file. Custom Crawlers are implicitly added.
+    CRAWLER_ARGS: {}
+    # Arguments for the DataReader's Dataset. The Dataset itself may be a custom Dataset, defined in a separate file. Custom Dataset are implicitly added.
+    DATASET_ARGS: 
+    # Whether to switch out the Generator in the DataReader with another built-in Generator. Note: custom Generators are implicitly added
+    GENERATOR: 
+    # Arguments for the DataReader's Generator. The Generator itself may be a custom Generator, defined in a separate file. Custom Generator are implicitly added.
+    GENERATOR_ARGS: 
+  # PLUGIN manages how plugins are used in this Deployment
+  PLUGIN:
+    # Controls when plugins are used. `always` means plugins are always used. `warmup` means plugins are used only until they are ready (i.e. activated)
+    # `activated` means only plugins that are ready are used.
+    HOOKS: always # always | warmup | activated
+    # Whether to reset already generated PLUGINS.
+
+# Manages the SAVE configuration of the EXECUTION or DEPLOYMENT, as well as file names. SAVE is essentially a readable unique identifier for a model
+# All entries in SAVE are purely for information purposes.
 SAVE:
+  # The version for this model
   MODEL_VERSION: 1
-  MODEL_CORE_NAME: "colabel_devel"
-  MODEL_BACKBONE: "res18"
+  # The name for this model
+  MODEL_CORE_NAME: "model"
+  # The backbone used for this model
+  MODEL_BACKBONE: "backbone"
+  # Any additional qualifiers for this model
   MODEL_QUALIFIER: "all"
+  # Whether to backup this model in a remote storage
   DRIVE_BACKUP: False
+  # Whether to backup this model's log files in a remote storage
+  LOG_BACKUP: False
+  # How often to perform a local save and backup for this model
   SAVE_FREQUENCY: 5 # Epoch
-  CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/CoLabeler/Models/"
+  # The URL where the backup will be conducted.
+  CHECKPOINT_DIRECTORY: "checkpoint"
 
+# TRANSFORMATION manages data transformation for both training and testing subsets
 TRANSFORMATION:
-  BATCH_SIZE: 8
-  WORKERS: 1
-  ARGS:
-    i_shape: [100,100]
-    normalization_mean: 0.5
-    normalization_std: 0.5
-    normalization_scale: 0.5
-    random_erase: False
-    random_erase_value: 0.5
-    channels: 3
+  # The number of samples in each batch
+  BATCH_SIZE: 32
+  # The number of workers to use in loading data
+  WORKERS: 2
+  # Arguments for The DATAREADER's Generator's transformations. See Generator's documentation for specifics.
+  ARGS: {}
 
+# TRAIN_TRANSFORMATION overwrites the TRANSFORMATION content for training data
 TRAIN_TRANSFORMATION:
-  ARGS:
-    h_flip: 0.5
-    t_crop: True
-
+  # The number of samples in each batch
+  BATCH_SIZE: 32
+  # The number of workers to use in loading data
+  WORKERS: 2
+  # Arguments for The DATAREADER's Generator's transformations. See Generator's documentation for specifics.
+  ARGS: {}
+  
+# TEST_TRANSFORMATION overwrites the TRANSFORMATION content for test data
 TEST_TRANSFORMATION:
-  ARGS:
-    h_flip: 0
+  # The number of samples in each batch
+  BATCH_SIZE: 32
+  # The number of workers to use in loading data
+  WORKERS: 2
+  # Arguments for The DATAREADER's Generator's transformations. See Generator's documentation for specifics.
+  ARGS: {}
 
+# MODEL manages model configuration
 MODEL:
-  BUILDER: classification_model_builder
-  MODEL_ARCH: 'ClassificationResnet'
-  MODEL_BASE: 'resnet18'
-  EMBEDDING_DIMENSIONS:
-  MODEL_NORMALIZATION: 'bn'
+  # Which model_builder to use. A custom model builder can be implicitly added. A model_builder simple generates a model from the provided arguments.
+  BUILDER: ednaml_model_builder
+  # What architecture to use for the model, e.g. HFAutoModel for a HuggingFace model, or a custom architecture. Architectures should inherit from ModelAbstract, 
+  # and use the remaining arguments.
+  MODEL_ARCH: ModelAbstract
+  # Argument for a ModelAbstract class that controls the feature extractor backbone, if used. For example, a MODEL_ARCH of ClassificationResnet accepts `resnet18` as a MODEL_BASE
+  MODEL_BASE: base
+  # The feature normalization, if used
+  MODEL_NORMALIZATION: bn
+  # Additional keyword arguments for the MODEL_ARCH provided. 
   MODEL_KWARGS:
+  # ModelAbstract models can be split into groups for more controlled optimization. A list of parameter groups can be provided here if your specific ModelAbstract, e.g. a GAN, supports it.
   PARAMETER_GROUPS: 
     - opt-1
 
-COLABEL:
-  COMPRESSION: 5
-  ENSEMBLE: "majority"
-  CONFIDENCE: "decay"
-
+# LOSS manages the losses used in the model training. Unused with EdnaDeploy. LOSS is a list of losses, 
+# one for each output or label of a model. LOSSes are indexed with their name. Each entry in LOSS can be a sum of multiple loss functions.
 LOSS:
-  - LOSSES: ['SoftmaxLogitsLoss']
+    # A list of LOSSES. Each entry is a string reference to an implemented Loss in EdnaML. PyTorch's built-in losses are accessed as TorchLoss.
+  - LOSSES: ['TorchLoss']
+    # Keyword arguments for each Loss.
     KWARGS: [{}]
+    # Weight for each loss
     LAMBDAS: [1.0]
+    # The label in the Crawler this Loss is targeting. Used when necessary
+    LABEL: 
+    # The name for this loss.
     NAME: out1
 
+# OPTIMIZER managers parameters for the training optimizer. Unused with EdnaDeploy. It contains a list of Optimizers, one for each parameter ground in a ModelAbstract
 OPTIMIZER:
-  OPTIMIZER_NAME: "Adam"
-  OPTIMIZER_KWARGS: '{}'
-  BASE_LR: 0.00001
-  LR_BIAS_FACTOR: 1.0
-  WEIGHT_DECAY: 0.0005
-  WEIGHT_BIAS_FACTOR: 0.0005
-  FP16: True
+    # Name of optimizer. This should match a parameter group in the ModelAbstract
+  - OPTIMIZER_NAME: opt-1
+    # Optimizer class. This should match a built-in Pytorch optimizer (see torch.optim), e.g. Adam, or Adagrad, etc
+    OPTIMIZER: Adam
+    # Arguments for the Optimizer class. 
+    OPTIMIZER_KWARGS: {}
+    # The Base learning rate for the Optimizer
+    BASE_LR: 0.00001
+    # The learning rate multiplier for bias terms in parameters
+    LR_BIAS_FACTOR: 1.0
+    # Weight decay factor, especially with Adam and variants
+    WEIGHT_DECAY: 0.0005
+    # Weight decay multiplier for bias terms in parameters
+    WEIGHT_BIAS_FACTOR: 0.0005
   
+# SCHEDULER managers parameters for optimizer learning rate scheduling
 SCHEDULER:
-  - LR_SCHEDULER: 'StepLR'
+    # A scheduler class, either from torch's built-in schedulers, or an EdnaML scheduler, such as FineGrainedSteppedLR
+  - LR_SCHEDULER: StepLR
+    # The name for this Scheduler. This should correspond to a parameter group in ModelAbstract, and match up with a corresponding Optimizer
+    SCHEDULER_NAME: opt-1
+    # Keyword arguments for the Scheduler
     LR_KWARGS: 
       step_size: 20
-      gamma: 0.1
 
+# MODEL_PLUGIN manages plugins that may be attached to a trained model.
+MODEL_PLUGIN:
+    # The name for this PLUGIN. This will be used as its identifier for this model
+  - PLUGIN_NAME: mp-1
+    # The PLUGIN class. Either a built-in plugin from ednaml.plugins or a Custom PLUGINS. For Custom PLUGINS, they must be stated here explicitly, NOT implicitly.
+    PLUGIN: ModelPlugin
+    # Keyword arguments for this PLUGIN
+    PLUGIN_KWARGS: {}
+
+
+# LOGGING manages any additional parameters, especially for logging
 LOGGING:
-  STEP_VERBOSE: 100 # Batch
\ No newline at end of file
+  # The number of training steps before printing intermediate losses and training accuracy.
+  STEP_VERBOSE: 100 # Batch
+  # The size for the input when estimating the Model's forward pass and getting the Model Summary. Should correspond with the expected input size.
+  INPUT_SIZE: Null
\ No newline at end of file
diff --git a/usage-docs/config.md b/usage-docs/config.md
index f8552077..394f2b71 100644
--- a/usage-docs/config.md
+++ b/usage-docs/config.md
@@ -20,7 +20,6 @@ The `EXECUTION` section contains parameters for EdnaML pipeline execution. They
     - `CRAWLER_ARGS`: A key-value dictionary or keyword arguments to pass into the datareader's crawler. Documentation is provided in [`datareaders.md`](config-reference/datareaders.md).
     - `DATASET_ARGS`: A key-value dictionary or keyword arguments to pass into the datareader's TorchDataset. Documentation is provided in [`datareaders.md`](config-reference/datareaders.md).
     - `GENERATOR_ARGS`: A key-value dictionary or keyword arguments to pass into the datareader's Generator. Documentation is provided in [`datareaders.md`](config-reference/datareaders.md).
-    - `DATALOADER_ARGS`: A key-value dictionary or keyword arguments to pass into the datareader's TorchDataloader. Documentation is provided in [`datareaders.md`](config-reference/datareaders.md).
 
 
 ## Sample
diff --git a/usage-docs/sample-configs/basics/cnn/cifar.yml b/usage-docs/sample-configs/basics/cnn/cifar.yml
new file mode 100644
index 00000000..1b02c434
--- /dev/null
+++ b/usage-docs/sample-configs/basics/cnn/cifar.yml
@@ -0,0 +1,89 @@
+EXECUTION:
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  MODEL_SERVING: False
+  EPOCHS: 5
+  SKIPEVAL: False
+  TEST_FREQUENCY: 1
+  FP16: False
+  DATAREADER: 
+    DATAREADER: TorchvisionDatareader
+    GENERATOR_ARGS:
+      tv_dataset: CIFAR10
+      tv_args: 
+        root: "Data/"
+        args:
+          download: true
+    DATASET_ARGS:
+      label_name: label
+      num_labels:
+  TRAINER: ClassificationTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: cifar10_resnet
+  MODEL_BACKBONE: res18
+  MODEL_QUALIFIER: cifar
+  DRIVE_BACKUP: False
+  LOG_BACKUP: False
+  SAVE_FREQUENCY: 5 # Epoch
+
+
+TRANSFORMATION:
+  BATCH_SIZE: 32
+  WORKERS: 2
+  ARGS:
+    i_shape: [32,32]
+    normalization_mean: 0.1307
+    normalization_std: 0.3081
+    normalization_scale: 0.5
+    random_erase: False
+    random_erase_value: 0.5
+    channels: 3
+
+TRAIN_TRANSFORMATION:
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+
+TEST_TRANSFORMATION:
+  ARGS:
+    h_flip: 0
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: ClassificationResnet
+  MODEL_BASE: resnet18
+  EMBEDDING_DIMENSIONS:
+  MODEL_NORMALIZATION: bn
+  MODEL_KWARGS:
+  PARAMETER_GROUPS: 
+    - opt-1
+
+LOSS:
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    NAME: out1
+    LABEL: label
+
+OPTIMIZER:
+  - OPTIMIZER: Adam
+    OPTIMIZER_NAME: opt-1
+    OPTIMIZER_KWARGS: {}
+    BASE_LR: 1.0e-3
+    LR_BIAS_FACTOR: 1.0
+    WEIGHT_DECAY: 0.0005
+    WEIGHT_BIAS_FACTOR: 0.0005
+  
+SCHEDULER:
+  - LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+    LR_KWARGS: 
+      step_size: 5
+      gamma: 0.1
+
+LOGGING:
+  STEP_VERBOSE: 100 # Batch
+  INPUT_SIZE: [32, 3, 32, 32] 
\ No newline at end of file
diff --git a/usage-docs/sample-configs/basics/cnn/cifar100.yml b/usage-docs/sample-configs/basics/cnn/cifar100.yml
new file mode 100644
index 00000000..a7cf5f32
--- /dev/null
+++ b/usage-docs/sample-configs/basics/cnn/cifar100.yml
@@ -0,0 +1,108 @@
+EXECUTION:
+  OPTIMIZER_BUILDER: ClassificationOptimizer
+  MODEL_SERVING: False
+  EPOCHS: 10
+  SKIPEVAL: True
+  TEST_FREQUENCY: 5
+  FP16: False
+  DATAREADER: 
+    DATAREADER: TorchvisionDatareader
+    GENERATOR_ARGS:
+      tv_dataset: CIFAR100
+      tv_args: 
+        root: "Data/"
+        args:
+          download: true
+    DATASET_ARGS:
+      label_name: label
+      num_labels:
+  TRAINER: ClassificationTrainer
+  TRAINER_ARGS: 
+    accumulation_steps: 1
+
+SAVE:
+  MODEL_VERSION: 1
+  MODEL_CORE_NAME: configtest
+  MODEL_BACKBONE: res18
+  MODEL_QUALIFIER: cifar100
+  DRIVE_BACKUP: False
+  SAVE_FREQUENCY: 5 # Epoch
+  CONFIG_BACKUP:
+    STORAGE: mlflow
+    FREQUENCY: 0  # i.e. only once
+  LOG_BACKUP:
+    STORAGE: mlflow
+    FREQUENCY: -1 # never
+  MODEL_BACKUP:
+    STORAGE: mlflow
+    FREQUENCY: # If none, use SAVE_FREQUENCY
+  MODEL_ARTIFACTS_BACKUP:
+    STORAGE: mlflow
+  MODEL_PLUGIN_BACKUP:
+    STORAGE: mlflow
+  METRICS_BACKUP:
+    STORAGE: mlflow
+
+STORAGE:
+  - NAME: mlflow
+    TYPE: MlflowStorage
+    STORAGE_ARGS:
+    URL: "./Backup/"
+
+TRANSFORMATION:
+  BATCH_SIZE: 32
+  WORKERS: 2
+  ARGS:
+    i_shape: [32,32]
+    normalization_mean: 0.1307
+    normalization_std: 0.3081
+    normalization_scale: 0.5
+    random_erase: False
+    random_erase_value: 0.5
+    channels: 3
+
+TRAIN_TRANSFORMATION:
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+
+TEST_TRANSFORMATION:
+  ARGS:
+    h_flip: 0
+
+MODEL:
+  BUILDER: ednaml_model_builder
+  MODEL_ARCH: ClassificationResnet
+  MODEL_BASE: resnet18
+  EMBEDDING_DIMENSIONS:
+  MODEL_NORMALIZATION: bn
+  MODEL_KWARGS:
+    
+  PARAMETER_GROUPS: 
+    - opt-1
+
+LOSS:
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    NAME: out1
+    LABEL: label
+
+OPTIMIZER:
+  - OPTIMIZER: Adam
+    OPTIMIZER_NAME: opt-1
+    OPTIMIZER_KWARGS: {}
+    BASE_LR: 1.0e-4
+    LR_BIAS_FACTOR: 1.0
+    WEIGHT_DECAY: 0.0005
+    WEIGHT_BIAS_FACTOR: 0.0005
+  
+SCHEDULER:
+  - LR_SCHEDULER: StepLR
+    SCHEDULER_NAME: opt-1
+    LR_KWARGS: 
+      step_size: 5
+      gamma: 0.1
+
+LOGGING:
+  STEP_VERBOSE: 100 # Batch
\ No newline at end of file
diff --git a/usage-docs/sample-configs/basics/cnn/mnist.yml b/usage-docs/sample-configs/basics/cnn/mnist.yml
new file mode 100644
index 00000000..473695d7
--- /dev/null
+++ b/usage-docs/sample-configs/basics/cnn/mnist.yml
@@ -0,0 +1,80 @@
+# Here, we set up the EXECUTION cell, which controls the EdnaML experiment itself
+EXECUTION:
+  EPOCHS: 5  # We will execute for 10 epochs
+  TEST_FREQUENCY: 1 # We will test the model every 1 epoch
+  DATAREADER: 
+    DATAREADER: TorchvisionDatareader   # We will use the built-in Torchvision Datareader to download and load MNIST
+    GENERATOR_ARGS:
+      tv_dataset: MNIST   # This tells our datareader that we want MNIST
+      tv_args: 
+        root: "Data/"     # This tells where to download the dataset
+        args:
+          download: true
+    DATASET_ARGS:
+      label_name: mnist_digits
+  TRAINER: ClassificationTrainer  # We will use a built-in trainer for Classification
+  TRAINER_ARGS: 
+    accumulation_steps: 2         # Number of steps for gradient accumulation
+
+
+SAVE:
+  MODEL_VERSION: 1              # We are only running v1 of this experiment
+  MODEL_CORE_NAME: mnist_resnet # The core name for this model
+  MODEL_BACKBONE: res18         # What backbone we are using here
+  MODEL_QUALIFIER: mnist        # Any other details we want to capture in the name
+  DRIVE_BACKUP: False           # Whether we are performing remote backup of the model
+  LOG_BACKUP: False             # Whether we are performing remote backup of the log files
+  SAVE_FREQUENCY: 5             # How often to save the model locally, and perform backups, if desired
+
+TRANSFORMATION:
+  BATCH_SIZE: 32                # The size of batches to provide to the model during training
+  WORKERS: 2                    # The number of workers to use during training
+  ARGS:                         # Additional args for this Generator (e.g. ednaml.generators.TorchvisionGeneratorWrapper)
+    i_shape: [28,28]
+    normalization_mean: 0.1307
+    normalization_std: 0.3081
+    normalization_scale: 0.5
+    random_erase: False
+    random_erase_value: 0.5
+    channels: 1
+
+TRAIN_TRANSFORMATION:           # Any replacements for the generator during training
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+
+TEST_TRANSFORMATION:            # Any replacements for the generator during testing
+  ARGS:
+    h_flip: 0
+
+MODEL:                              # These are model-specific details
+  BUILDER: ednaml_model_builder     # ednaml_model_builder is a basic model_builder that verifies model arguments make sense
+  MODEL_ARCH: ClassificationResnet  # The built-in architecture we are using
+  MODEL_BASE: resnet18              # The base we are using for the architecture. For example, ClassificationResnet accepts resnet18, resnet34, resnet50, etc
+  MODEL_KWARGS:
+    initial_channels: 1             # Since MNIST is black and white with 1 channel, we use the `initial_channels` parameter of ClassificationResnet to set this
+  
+LOSS:
+  - LOSSES: ['SoftmaxLogitsLoss']   # We use the basic softmax loss here
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    NAME: digit_out                 # The name for this loss
+    LABEL: mnist_digits             # The label this loss is targeting. This MUST match the EXECUTION.DATAREADER.DATASET_ARGS["label_name"] above
+
+OPTIMIZER:
+  - OPTIMIZER: Adam                 # We use the standard Adam optimizer
+    OPTIMIZER_KWARGS: {}
+    BASE_LR: 1.0e-3
+    LR_BIAS_FACTOR: 1.0
+    WEIGHT_DECAY: 0.0005
+    WEIGHT_BIAS_FACTOR: 0.0005
+  
+SCHEDULER:
+  - LR_SCHEDULER: StepLR            # The basic optimizer scheduler to adynamically adjust learning rate
+    LR_KWARGS: 
+      step_size: 5
+      gamma: 0.1
+
+LOGGING:
+  STEP_VERBOSE: 100                 # How often to print intermediate results
+  INPUT_SIZE: [32, 1, 28, 28]       # The estimated shape of the input. 
\ No newline at end of file
diff --git a/usage-docs/sample-configs/trainers/multibranch-compcars.yml b/usage-docs/sample-configs/trainers/multibranch-compcars.yml
index b59e52ed..e7ba0fd2 100644
--- a/usage-docs/sample-configs/trainers/multibranch-compcars.yml
+++ b/usage-docs/sample-configs/trainers/multibranch-compcars.yml
@@ -8,6 +8,9 @@ EXECUTION:
     DATAREADER: CompCars
     CRAWLER_ARGS:
       data_folder: "Data/CompCars"
+      train_folder: "image"
+      trainfile: "train.txt"
+      testfile: "test.txt"
     DATASET_ARGS: 
       pathidx: 4
       annotationidx: [0,2,3]
@@ -28,18 +31,23 @@ SAVE:
   
 
 TRANSFORMATION:
-  SHAPE: [200,200]
-  NORMALIZATION_MEAN: 0.5
-  NORMALIZATION_STD: 0.5
-  NORMALIZATION_SCALE: 255
-  H_FLIP: 0.5
-  T_CROP: True
-  RANDOM_ERASE: False
-  RANDOM_ERASE_VALUE: 0.5
-  CHANNELS: 3
+  ARGS:
+    i_shape: [200,200]
+    normalization_mean: 0.5
+    normalization_std: 0.5
+    normalization_scale: 0.5
+    channels: 3
   BATCH_SIZE: 32
   WORKERS: 2
 
+TRAIN_TRANSFORMATION:
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+    random_erase: True
+    random_erase_value: 0.3
+
+
 MODEL:
   BUILDER: multibranch_model_builder
   MODEL_ARCH: 'MultiBranchResnet'
@@ -48,28 +56,33 @@ MODEL:
   MODEL_KWARGS: 
     number_branches: 3
     branches:
-      - name: make
+      - name: bmake
         number_outputs: 1
         outputs: 
           - dimensions:
             name: make-fc
             label: make
-      - name: type
+      - name: btype
         number_outputs: 1
         outputs: 
           - dimensions:
             name: type-fc
             label: type
-      - name: year
+      - name: byear
         number_outputs: 1
         outputs: 
           - dimensions:
             name: year-fc
             label: year
     fuse: True
-    fuse_outputs: ['make', 'year', 'type']
+    fuse_outputs: ['bmake', 'byear', 'btype']
     fuse_dimensions:
     fuse_label: make
+    fuse_name: fuse
+    shared_block: 1
+    soft_targets: True
+    soft_target_branch: ['bmake']
+    soft_target_output_source: fuse
     attention: 'cbam'
 
 COLABEL:
@@ -81,24 +94,43 @@ LOSS:
   - LOSSES: ['SoftmaxLogitsLoss', 'SoftmaxLabelSmooth']
     KWARGS: [{}, {'eps':0.2}]
     LAMBDAS: [1.0, 1.0]
-    NAME: color
+    LABEL: year
+    NAME: year-fc
   - LOSSES: ['SoftmaxLabelSmooth', 'SoftmaxLogitsLoss']
     KWARGS: [{'eps':0.3}, {}]
     LAMBDAS: [1.0, 2.0]
-    NAME: type
+    LABEL: type
+    NAME: type-fc
+  - LOSSES: ['SoftmaxLogitsLoss', 'SoftmaxLabelSmooth']
+    KWARGS: [{}, {'eps':0.2}]
+    LAMBDAS: [1.0, 1.0]
+    LABEL: make
+    NAME: make-fc
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    LABEL: make
+    NAME: fuse
+  - LOSSES: ['SoftmaxLogitsLoss']
+    KWARGS: [{}]
+    LAMBDAS: [1.0]
+    LABEL: make
+    NAME: bmake # This is for soft targets
+  
 
 OPTIMIZER:
-  OPTIMIZER_NAME: "Adam"
-  OPTIMIZER_KWARGS: '{}'
-  BASE_LR: 0.0001
-  LR_BIAS_FACTOR: 1.0
-  WEIGHT_DECAY: 0.0005
-  WEIGHT_BIAS_FACTOR: 0.0005
-  FP16: True
+  - OPTIMIZER: "Adam"
+    OPTIMIZER_NAME: opt-1
+    OPTIMIZER_KWARGS: {}
+    BASE_LR: 0.0001
+    LR_BIAS_FACTOR: 1.0
+    WEIGHT_DECAY: 0.0005
+    WEIGHT_BIAS_FACTOR: 0.0005
+    FP16: True
 
 SCHEDULER:
-  LR_SCHEDULER: 'StepLR'
-  LR_KWARGS: '{"step_size":20, "gamma":0.1}'
+  - LR_SCHEDULER: 'StepLR'
+    LR_KWARGS: {"step_size":20, "gamma":0.1}
 
 LOGGING:
   STEP_VERBOSE: 100 # Batch
\ No newline at end of file
diff --git a/usage-docs/sample-configs/trainers/multibranch-veri.yml b/usage-docs/sample-configs/trainers/multibranch-veri.yml
index b73fb0a3..15ade4ba 100644
--- a/usage-docs/sample-configs/trainers/multibranch-veri.yml
+++ b/usage-docs/sample-configs/trainers/multibranch-veri.yml
@@ -28,18 +28,23 @@ SAVE:
   
 
 TRANSFORMATION:
-  SHAPE: [200,200]
-  NORMALIZATION_MEAN: 0.5
-  NORMALIZATION_STD: 0.5
-  NORMALIZATION_SCALE: 255
-  H_FLIP: 0.5
-  T_CROP: True
-  RANDOM_ERASE: False
-  RANDOM_ERASE_VALUE: 0.5
-  CHANNELS: 3
+  ARGS:
+    i_shape: [200,200]
+    normalization_mean: 0.5
+    normalization_std: 0.5
+    normalization_scale: 0.5
+    channels: 3
   BATCH_SIZE: 64
   WORKERS: 2
 
+TRAIN_TRANSFORMATION:
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+    random_erase: True
+    random_erase_value: 0.3
+
+
 MODEL:
   BUILDER: multibranch_model_builder
   MODEL_ARCH: 'MultiBranchResnet'
diff --git a/usage-docs/sample-configs/trainers/multiclassification-veri.yml b/usage-docs/sample-configs/trainers/multiclassification-veri.yml
index dfda0477..9b67b32e 100644
--- a/usage-docs/sample-configs/trainers/multiclassification-veri.yml
+++ b/usage-docs/sample-configs/trainers/multiclassification-veri.yml
@@ -26,20 +26,23 @@ SAVE:
   SAVE_FREQUENCY: 5 # Epoch
   CHECKPOINT_DIRECTORY: "./drive/MyDrive/Projects/CoLabeler/Models/"
   
-
 TRANSFORMATION:
-  SHAPE: [200,200]
-  NORMALIZATION_MEAN: 0.5
-  NORMALIZATION_STD: 0.5
-  NORMALIZATION_SCALE: 255
-  H_FLIP: 0.5
-  T_CROP: True
-  RANDOM_ERASE: False
-  RANDOM_ERASE_VALUE: 0.5
-  CHANNELS: 3
-  BATCH_SIZE: 128
+  ARGS:
+    shape: [200,200]
+    normalization_mean: 0.5
+    normalization_std: 0.5
+    normalization_scale: 0.5
+    channels: 3
+  BATCH_SIZE: 64
   WORKERS: 2
 
+TRAIN_TRANSFORMATION:
+  ARGS:
+    h_flip: 0.5
+    t_crop: True
+    random_erase: True
+    random_erase_value: 0.3
+
 MODEL:
   BUILDER: multiclassification_model_builder
   MODEL_ARCH: 'MultiClassificationResnet'
diff --git a/samples/EdnaML Custom Crawler.ipynb b/usage-docs/sample-notebooks/EdnaML Custom Crawler.ipynb
similarity index 100%
rename from samples/EdnaML Custom Crawler.ipynb
rename to usage-docs/sample-notebooks/EdnaML Custom Crawler.ipynb
diff --git a/samples/EdnaML Debugging.ipynb b/usage-docs/sample-notebooks/EdnaML Debugging.ipynb
similarity index 100%
rename from samples/EdnaML Debugging.ipynb
rename to usage-docs/sample-notebooks/EdnaML Debugging.ipynb


Test patch:
diff --git a/tests/test_config_merge.py b/tests/test_config_merge.py
new file mode 100644
index 00000000..ec400691
--- /dev/null
+++ b/tests/test_config_merge.py
@@ -0,0 +1,121 @@
+# Testing merging of configs
+
+cfg_base = \
+"""
+A:
+  B: alpha
+  C: beta
+  D:
+    e: delta
+    f: gamma
+  G:
+    h: epsilon
+    i: upsilon
+  J:
+    K: omicron
+    L: 
+      M: phi
+      N:
+        o: psi
+        p: yotta
+Q:
+  R: zeta
+  S: tau
+  T:
+    u: pi
+    v: sigma
+"""
+
+cfg_extension_a = \
+"""
+A:
+  B: modified_alpha
+  C: modified_beta
+
+W:
+  X: chi
+  Y: eta
+  Z:
+    a: mu
+    b: nu
+"""
+
+cfg_extension_b = \
+"""
+A:
+  D:
+  G:
+    h: modified_epsilon
+  J: deleted_key
+
+Q:
+  S: modified_tau
+  T:
+    u: sigma
+"""
+
+cfg_base_plus_a = \
+"""
+A:
+  B: modified_alpha
+  C: modified_beta
+  D:
+    e: delta
+    f: gamma
+  G:
+    h: epsilon
+    i: upsilon
+  J:
+    K: omicron
+    L: 
+      M: phi
+      N:
+        o: psi
+        p: yotta
+Q:
+  R: zeta
+  S: tau
+  T:
+    u: pi
+    v: sigma
+
+W:
+  X: chi
+  Y: eta
+  Z:
+    a: mu
+    b: nu
+"""
+
+cfg_base_plus_b = \
+"""
+A:
+  B: alpha
+  C: beta
+  D:
+  G:
+    h: modified_epsilon
+  J: deleted_key
+
+Q:
+  R: zeta
+  S: modified_tau
+  T:
+    u: sigma
+"""
+
+
+def test_merge_a():
+    test_file=  "base_test.yml"
+    test_file_a=  "test_a.yml"
+
+    with open(test_file, "w") as wfile:
+        wfile.write(cfg_base)
+
+    with open(test_file_a, "w") as wfile:
+        wfile.write(cfg_extension_a)
+    
+    
+    import ednaml
+    from ednaml.core import EdnaML
+    eml = EdnaML()
\ No newline at end of file