--- a
+++ b/medacy/data/dataset.py
@@ -0,0 +1,302 @@
+"""
+A medaCy Dataset facilities the management of data for both model training and model prediction.
+
+A Dataset object provides a wrapper for a unix file directory containing training/prediction
+data. If a Dataset, at training time, is fed into a pipeline requiring auxilary files
+(Metamap for instance) the Dataset will automatically create those files in the most efficient way possible.
+
+Training
+#################
+When a directory contains **both** raw text files alongside annotation files, an instantiated Dataset
+detects and facilitates access to those files.
+
+Assuming your directory looks like this (where .ann files are in `BRAT <http://brat.nlplab.org/standoff.html>`_ format):
+::
+    home/medacy/data
+    ├── file_one.ann
+    ├── file_one.txt
+    ├── file_two.ann
+    └── file_two.txt
+
+A common data work flow might look as follows.
+
+Running:
+::
+    >>> from medacy.data import Dataset
+    >>> from medacy.pipeline_components.feature_overlayers.metamap.metamap import MetaMap
+
+    >>> dataset = Dataset('/home/datasets/some_dataset')
+    >>> for data_file in dataset:
+    ...    (data_file.file_name, data_file.raw_path, dataset.ann_path)
+    (file_one, file_one.txt, file_one.ann)
+    (file_two, file_two.txt, file_two.ann)
+    >>> dataset
+    ['file_one', 'file_two']
+    >>>> dataset.is_metamapped()
+    False
+    >>> metamap = MetaMap('/home/path/to/metamap/binary')
+    >>> with metamap:
+    ...     metamap.metamap_dataset(dataset)
+    >>> dataset.is_metamapped()
+    True
+
+MedaCy **does not** alter the data you load in any way - it only reads from it.
+
+Prediction
+##########
+When a directory contains **only** raw text files, an instantiated Dataset object interprets this as
+a directory of files that need to be predicted. This means that the internal Datafile that aggregates
+meta-data for a given prediction file does not have fields for annotation_file_path set.
+
+When a directory contains **only** ann files, an instantiated Dataset object interprets this as
+a directory of files that are predictions. Useful methods for analysis include :meth:`medacy.data.dataset.Dataset.compute_confusion_matrix`,
+:meth:`medacy.data.dataset.Dataset.compute_ambiguity` and :meth:`medacy.data.dataset.Dataset.compute_counts`.
+
+External Datasets
+#################
+
+In the real world, datasets (regardless of domain) are evolving entities. Hence, it is essential to version them.
+A medaCy compatible dataset can be created to facilitate this versioning. A medaCy compatible dataset lives a python
+packages that can be hooked into medaCy or used for any other purpose - it is simply a loose wrapper for this Dataset
+object. Instructions for creating such a dataset can be found `here <https://github.com/NLPatVCU/medaCy/tree/master/examples/guide>`_.
+wrap them.
+"""
+
+import argparse
+import json
+import logging
+import os
+import pprint
+from collections import Counter
+from pathlib import Path
+
+from medacy.data.annotations import Annotations
+from medacy.data.data_file import DataFile
+
+
+class Dataset:
+    """
+    A facilitation class for data management.
+    """
+
+    def __init__(self, data_directory, data_limit=None):
+        """
+        Manages directory of training data along with other medaCy generated files.
+
+        Only text files: considers a directory for managing metamapping.
+        Only ann files: considers a directory of predictions.
+        Both text and ann files: considers a directory for training.
+
+        :param data_directory: Directory containing data for training or prediction.
+        :param data_limit: A limit to the number of files to process. Must be between 1 and number of raw text files in data_directory
+        """
+        self.data_directory = Path(data_directory)
+
+        metamap_dir = self.data_directory / 'metamapped'
+        self.metamapped_files_directory = metamap_dir if metamap_dir.is_dir() else None
+
+        self.data_files = self._create_data_files()
+        self.data_limit = data_limit or len(self.data_files)
+
+    def _create_data_files(self):
+        data_files = []
+        all_files_in_directory = os.listdir(self.data_directory)
+        all_file_base_names = {f.split(".")[0] for f in all_files_in_directory}
+
+        for file_name in all_file_base_names:
+            txt_path = None
+            ann_path = None
+            metamapped_path = None
+
+            potential_txt_path = self.data_directory / (file_name + ".txt")
+            if potential_txt_path.exists():
+                txt_path = potential_txt_path
+
+            potential_ann_path = self.data_directory / (file_name + ".ann")
+            if potential_ann_path.exists():
+                ann_path = potential_ann_path
+
+            if self.metamapped_files_directory:
+                potential_mm_path = self.metamapped_files_directory / (file_name + ".metamapped")
+                if potential_mm_path.exists():
+                    metamapped_path = potential_mm_path
+
+            if txt_path or ann_path:
+                new_df = DataFile(file_name, txt_path, ann_path, metamapped_path)
+                data_files.append(new_df)
+
+        return sorted(data_files, key=lambda x: x.file_name)
+
+    def __iter__(self):
+        return iter(self.data_files[0:self.data_limit])
+
+    def __len__(self):
+        return len(self.data_files)
+
+    def is_metamapped(self):
+        """
+        Verifies if all fil es in the Dataset are metamapped.
+
+        :return: True if all data files are metamapped, False otherwise.
+        """
+        if self.metamapped_files_directory is None or not self.metamapped_files_directory.exists():
+            return False
+
+        for file in self.data_files:
+            potential_file_path = self.metamapped_files_directory / f"{file.file_name}.metamapped"
+            if not potential_file_path.exists():
+                return False
+
+            # Metamapped file could exist, but metamapping it could have failed.
+            # If the file is less than 200 bytes, log a warning.
+            file_size_in_bytes = os.path.getsize(potential_file_path)
+            if file_size_in_bytes < 200:
+                logging.warning(f"Metamapped version of {file.file_name} is only {file_size_in_bytes} bytes. "
+                                f"Metamapping could have failed: {potential_file_path}")
+
+        return True
+
+    def __str__(self):
+        """
+        Prints a list-like string of the names of the Datafile objects up to the data limit
+        (can't be used if copied and pasted)
+        """
+        return str([d.file_name for d in self])
+
+    def compute_counts(self):
+        """
+        Computes entity counts over all documents in this dataset.
+
+        :return: a Counter of entity counts
+        """
+        total = Counter()
+
+        for ann in self.generate_annotations():
+            total += ann.compute_counts()
+
+        return total
+
+    def compute_confusion_matrix(self, other, leniency=0):
+        """
+        Generates a confusion matrix where this Dataset serves as the gold standard annotations and `dataset` serves
+        as the predicted annotations. A typical workflow would involve creating a Dataset object with the prediction directory
+        outputted by a model and then passing it into this method.
+
+        :param other: a Dataset object containing a predicted version of this dataset.
+        :param leniency: a floating point value between [0,1] defining the leniency of the character spans to count as different. A value of zero considers only exact character matches while a positive value considers entities that differ by up to :code:`ceil(leniency * len(span)/2)` on either side.
+        :return: two element tuple containing a label array (of entity names) and a matrix where rows are gold labels and columns are predicted labels. matrix[i][j] indicates that entities[i] in this dataset was predicted as entities[j] in 'annotation' matrix[i][j] times
+        """
+        if not isinstance(other, Dataset):
+            raise ValueError("other must be instance of Dataset")
+
+        # verify files are consistent
+        diff = {d.file_name for d in self} - {d.file_name for d in other}
+        if diff:
+            raise ValueError(f"Dataset of predictions is missing the files: {repr(diff)}")
+
+        # sort entities in ascending order by count.
+        entities = [key for key, _ in sorted(self.compute_counts().items(), key=lambda x: x[1])]
+        confusion_matrix = [[0 * len(entities)] * len(entities)]
+
+        for gold_data_file in self:
+            prediction_iter = iter(other)
+            prediction_data_file = next(prediction_iter)
+            while str(gold_data_file) != str(prediction_data_file):
+                prediction_data_file = next(prediction_iter)
+
+            gold_annotation = Annotations(gold_data_file.ann_path)
+            pred_annotation = Annotations(prediction_data_file.ann_path)
+
+            # compute matrix on the Annotation file level
+            ann_confusion_matrix = gold_annotation.compute_confusion_matrix(pred_annotation, entities, leniency=leniency)
+            for i in range(len(confusion_matrix)):
+                for j in range(len(confusion_matrix)):
+                    confusion_matrix[i][j] += ann_confusion_matrix[i][j]
+
+        return entities, confusion_matrix
+
+    def compute_ambiguity(self, dataset):
+        """
+        Finds occurrences of spans from 'dataset' that intersect with a span from this annotation but do not have this spans label.
+        label. If 'dataset' comprises a models predictions, this method provides a strong indicators
+        of a model's in-ability to dis-ambiguate between entities. For a full analysis, compute a confusion matrix.
+
+        :param dataset: a Dataset object containing a predicted version of this dataset.
+        :return: a dictionary containing the ambiguity computations on each gold, predicted file pair
+        """
+        if not isinstance(dataset, Dataset):
+            raise ValueError("dataset must be instance of Dataset")
+
+        # verify files are consistent
+        diff = {d.file_name for d in self} - {d.file_name for d in dataset}
+        if diff:
+            raise ValueError(f"Dataset of predictions is missing the files: {repr(diff)}")
+
+        # Dictionary storing ambiguity over dataset
+        ambiguity_dict = {}
+
+        for gold_data_file in self:
+            prediction_iter = iter(dataset)
+            prediction_data_file = next(prediction_iter)
+            while str(gold_data_file) != str(prediction_data_file):
+                prediction_data_file = next(prediction_iter)
+
+            gold_annotation = Annotations(gold_data_file.ann_path)
+            pred_annotation = Annotations(prediction_data_file.ann_path)
+
+            # compute matrix on the Annotation file level
+            ambiguity_dict[str(gold_data_file)] = gold_annotation.compute_ambiguity(pred_annotation)
+
+        return ambiguity_dict
+
+    def get_labels(self, as_list=False):
+        """
+        Get all of the entities/labels used in the dataset.
+        :param as_list: bool for if to return the results as a list; defaults to False
+        :return: A set of strings. Each string is a label used.
+        """
+        labels = set()
+
+        for ann in self.generate_annotations():
+            labels.update(ann.get_labels())
+
+        if as_list:
+            return list(labels)
+        return labels
+
+    def generate_annotations(self):
+        """Generates Annotation objects for all the files in this Dataset"""
+        for file in self:
+            if file.ann_path is not None:
+                yield Annotations(file.ann_path, source_text_path=file.txt_path)
+            else:
+                yield Annotations([])
+
+    def __getitem__(self, item):
+        """
+        Creates and returns the Annotations object with the given file name, else raises FileNotFoundError;
+        useful for getting Annotations objects from parallel Datasets
+        :param item: the name of the file to be represented (not including the extension or parent directories)
+        :return: an Annotations object
+        """
+        path = os.path.join(self.data_directory, item + '.ann')
+        return Annotations(path)
+
+
+def main():
+    """CLI for retrieving dataset information"""
+    parser = argparse.ArgumentParser(description='Calculate data about a given data directory')
+    parser.add_argument('directory')
+    args = parser.parse_args()
+
+    dataset = Dataset(args.directory)
+
+    entities = json.dumps(dataset.get_labels(as_list=True))
+    counts = dataset.compute_counts()
+
+    print(f"Entities: {entities}")
+    pprint.pprint(counts)
+
+
+if __name__ == '__main__':
+    main()