Repository: asuprem/EdnaML (language: python)

Issue:
Suprem devel edeploy merge



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/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/MiDAS/midas-dataset.py b/profiles/MiDAS/midas-dataset.py
index 355afbd1..da67eae0 100644
--- a/profiles/MiDAS/midas-dataset.py
+++ b/profiles/MiDAS/midas-dataset.py
@@ -196,8 +196,8 @@ def build_transforms(self, transform, mode, **kwargs):  #<-- generator kwargs:
     self.tokenizer = self.tokenizer(**kwargs) # vocab_file, do_lower_case, spm_model_file
 
   
-  def buildDataset(self, crawler, mode, transform, **kwargs): #<-- dataset args:
-    return MiDASDataset(crawler.metadata[mode]["crawl"], mode, tokenizer = self.tokenizer, **kwargs) # needs maxlen, memcache, mlm_probability
+  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,
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/setup.cfg b/setup.cfg
index 405476a5..74217f9b 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = ednaml
-version=0.1.4
+version=0.1.5
 url = https://github.com/asuprem/GLAMOR
 project_urls =
     Documentation = https://github.com/asuprem/GLAMOR
@@ -37,6 +37,8 @@ azure =
     azure-storage-blob >= 12.13.0
 metrics = 
     torchmetrics >= 0.9.3
+transformers = 
+    transformers >= 4.21.*
 full = ednaml[azure,metrics]
 
 
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 3a43a7ec..b1c123b3 100644
--- a/src/ednaml/config/ConfigDefaults.py
+++ b/src/ednaml/config/ConfigDefaults.py
@@ -21,7 +21,7 @@ class ConfigDefaults(BaseConfig):
     CHECKPOINT_DIRECTORY: str
 
     OUTPUT_ARGS: Dict[str,str]
-    DEPLOYMENT: str
+    DEPLOY: str
     DEPLOYMENT_ARGS: Dict[str,str]
     DEPLOYMENT_EPOCHS: int
 
@@ -55,6 +55,7 @@ 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]
@@ -71,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",
@@ -82,17 +83,17 @@ 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)
 
@@ -127,6 +128,7 @@ 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", {})
diff --git a/src/ednaml/config/DeploymentConfig.py b/src/ednaml/config/DeploymentConfig.py
index 54c6dbdc..7300da94 100644
--- a/src/ednaml/config/DeploymentConfig.py
+++ b/src/ednaml/config/DeploymentConfig.py
@@ -20,8 +20,8 @@ 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", {})
diff --git a/src/ednaml/config/EdnaMLConfig.py b/src/ednaml/config/EdnaMLConfig.py
index a851a6c3..777b3799 100644
--- a/src/ednaml/config/EdnaMLConfig.py
+++ b/src/ednaml/config/EdnaMLConfig.py
@@ -3,6 +3,7 @@
 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
@@ -46,13 +47,39 @@ class EdnaMLConfig(BaseConfig):
     extensions: List[str]
 
     def __init__(
-        self, config_path: str, defaults: ConfigDefaults = ConfigDefaults()
+        self, config_path: List[str], defaults: ConfigDefaults = ConfigDefaults(), **kwargs
     ):
         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.read_path(config_path)
-
+        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
+
+
     def _has_extension_verifier(self, ydict, extension, verification):
         return False if ydict.get(extension, verification) == verification else True
 
@@ -174,7 +201,7 @@ 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:
@@ -185,14 +212,17 @@ def extend(self, config_path: str, defaults: ConfigDefaults = ConfigDefaults()):
             config_path = [config_path]
         if type(config_path) is list:
             for cpath in config_path:
-                responses[cpath] = self._extend(cpath, defaults)
+                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):
+    def _extend(self, config_path, defaults: ConfigDefaults, **kwargs):
         ydict = self.read_path(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)
         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])
 
@@ -222,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/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
index ec489a6e..3943749f 100644
--- a/src/ednaml/config/ExecutionPluginConfig.py
+++ b/src/ednaml/config/ExecutionPluginConfig.py
@@ -10,5 +10,5 @@ class ExecutionPluginConfig(BaseConfig):
 
     def __init__(self, ep_dict):
         #self.EPOCHS = datareader_dict.get("PLUGIN_EPOCHS", 1)
-        self.HOOKS = ep_dict.get("PLUGIN_HOOKS", 'always')
-        self.RESET = ep_dict.get("PLUGIN_RESET", False)
\ No newline at end of file
+        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/core/EdnaDeploy.py b/src/ednaml/core/EdnaDeploy.py
index be64199f..c5a360fa 100644
--- a/src/ednaml/core/EdnaDeploy.py
+++ b/src/ednaml/core/EdnaDeploy.py
@@ -105,7 +105,8 @@ 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,
@@ -189,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 8bcb14af..7b20b72d 100644
--- a/src/ednaml/core/EdnaML.py
+++ b/src/ednaml/core/EdnaML.py
@@ -82,23 +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])
-            for cfg_item in config[1:]:
-                msg = self.cfg.extend(cfg_item)
-                print(str(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...
@@ -176,6 +184,7 @@ def resetStorageQueues(self):
 
     def resetQueues(self):
         """Resets the `apply()` queue"""
+        self.logger.debug("Resetting declarative queues.")
         for queue_function in self.resetQueueArray:
             queue_function()
 
@@ -252,6 +261,9 @@ def apply(self, **kwargs):
         self.buildModel()
         self.loadWeights()
         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()
@@ -266,10 +278,12 @@ def apply(self, **kwargs):
         self.resetQueues()
 
     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
 
@@ -308,6 +322,7 @@ 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
 
@@ -728,6 +743,7 @@ def addPlugins(self, plugin_class_list: List[Type[ModelPlugin]], **kwargs):
         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()
@@ -758,10 +774,10 @@ def addPluginsToModel(self):
                 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)
+        #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
@@ -777,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
@@ -845,7 +862,10 @@ 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
         try:
@@ -888,6 +908,19 @@ def _add(self, file_or_module_path):
         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]
                 )
@@ -1061,7 +1094,7 @@ def buildTrainDataloader(
                 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
@@ -1263,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/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 3db93d45..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
@@ -86,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:
@@ -97,7 +101,6 @@ def apply(
     def saveMetadata(self):
         print("NOT saving metadata. saveMetadata() function not set up.")
 
-    
     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")
@@ -139,6 +142,9 @@ def deploy(self, continue_epoch=0, inference = False, ignore_plugins: List[str]
                 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`.")
@@ -148,13 +154,16 @@ def data_step(self):
             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    # TODO move plugins here to allow labels as well!!!!!!!!
@@ -162,7 +171,11 @@ def deploy_step(self, batch):   # USER IMPLEMENTS
 
         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")
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 8c2e20c4..8107385b 100644
--- a/src/ednaml/generators/AlbertGenerator.py
+++ b/src/ednaml/generators/AlbertGenerator.py
@@ -31,7 +31,9 @@ def __init__(self, logger, dataset, mode, transform=None, **kwargs):
             else:
                 self.shards_exist = True
                 self.logger.debug("Shards already exist and `shard_replace` is False")
-        os.makedirs(self.shardpath, exist_ok=True)
+        if self.shardcache:
+            self.logger.debug("Creating shardpath %s"%self.shardpath)
+            os.makedirs(self.shardpath, exist_ok=True)
         
 
 
@@ -119,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/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 332ffcad..8014b7b7 100644
--- a/src/ednaml/models/ModelAbstract.py
+++ b/src/ednaml/models/ModelAbstract.py
@@ -148,7 +148,7 @@ def forward(self, x, labels=None, **kwargs):    # TODO Labels for the 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) -> Tuple[TensorType,TensorType,List[Any]]:
+    def foward_impl(self, x, **kwargs) -> Tuple[TensorType,TensorType,List[Any]]:   # Return Logits, Features, Secondary Outputs
 
         raise NotImplementedError()
 
@@ -221,6 +221,7 @@ def addPlugin(self, plugin: Type[ModelPlugin], plugin_name: str = None, plugin_k
         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
         
@@ -235,20 +236,20 @@ def set_usable_plugins(self):
         elif self.plugin_hook == "activated":
             self.usable_plugins = [item for item in self.plugins if self.plugins[item].activated]
         else:
-            raise ValueError("Unknown calue for `plugin_hook`: %s"%self.plugin_hook)
+            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.plugins:
+        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.plugins:
+        for plugin in self.usable_plugins:
             self.plugins[plugin].post_epoch(model = self, epoch=epoch)
         self.set_usable_plugins()   # post plugin usage
 
@@ -256,7 +257,7 @@ 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.plugins:
+        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
@@ -265,7 +266,7 @@ def post_forward_hook(self, x, feature_logits, features, secondary_outputs, **kw
         # 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.plugins:
+        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/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
index 22b45b9e..30797c53 100644
--- a/src/ednaml/plugins/KMeansProxy.py
+++ b/src/ednaml/plugins/KMeansProxy.py
@@ -6,7 +6,7 @@
 
 class KMeansProxy(ModelPlugin):
     name = "KMeansProxy"
-    def __init__(self, num_clusters=10, dimensions=768, dist="euclidean", rand_seed=145845, epochs=3):
+    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)
 
 
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/RandomizedLipschitz.py b/src/ednaml/plugins/RandomizedLipschitz.py
index 02fea77a..872cf0c5 100644
--- a/src/ednaml/plugins/RandomizedLipschitz.py
+++ b/src/ednaml/plugins/RandomizedLipschitz.py
@@ -1,5 +1,4 @@
 
-import pdb
 from typing import List
 from ednaml.models.ModelAbstract import ModelAbstract
 from ednaml.plugins import ModelPlugin
@@ -12,7 +11,7 @@
 
 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):
+    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)
 
 
@@ -131,6 +130,10 @@ def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
             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()
@@ -144,15 +147,20 @@ def post_epoch(self, model: ModelAbstract, epoch: int = 0, **kwargs):
 
                         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.lipschitz_threshold = torch.max(l_scores).item()
-                        self.lipschitz_threshold_mean = torch.mean(l_scores).item()
-                        self.smooth_lipschitz_threshold = torch.max(smooth_lscores).item()
-                        self.smooth_lipschitz_threshold_mean = torch.mean(smooth_lscores).item()
-
-                        self.epsilon = 1. / self.smooth_lipschitz_threshold
+                self.epsilon = 1. / self.smooth_lipschitz_threshold
 
-                        self.activated = True
+                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
@@ -205,15 +213,6 @@ def search_minibatch_kmeans_clusters(self, batch: torch.Tensor, ):
         
         # For each cluster center, we have a dict of closest features...
 
-    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])
-
     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)
diff --git a/src/ednaml/trainer/BaseTrainer.py b/src/ednaml/trainer/BaseTrainer.py
index cd51a47a..e56c6660 100644
--- a/src/ednaml/trainer/BaseTrainer.py
+++ b/src/ednaml/trainer/BaseTrainer.py
@@ -1,6 +1,6 @@
 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
@@ -118,6 +118,10 @@ 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)
@@ -158,9 +162,14 @@ 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:
@@ -390,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()
@@ -442,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/__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 6571bc4f..c5ad8d60 100644
--- a/src/ednaml/utils/__init__.py
+++ b/src/ednaml/utils/__init__.py
@@ -8,15 +8,18 @@
 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)
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/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/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