setup

fake_job_detector/cli.py

@@ -0,0 +1,43 @@
+import argparse
+
+
+# Define the models we can choose from
+class ModelEnum:
+    SVM = "svm"
+    DISTILBERT = "distilbert"
+
+    @classmethod
+    def choices(cls):
+        return [cls.SVM, cls.DISTILBERT]
+
+# Define the CLI parser
+parser = argparse.ArgumentParser(description="CLI for predicting if a job is fake based on the title and description")
+subparsers = parser.add_subparsers(title="subcommands", dest="subcommand")
+
+# Prediction sub-command
+predict_parser = subparsers.add_parser("predict", help="Make predictions using a trained model")
+predict_parser.add_argument("-m", "--model", choices=ModelEnum.choices(), required=True, help="Model to choose between SVM baseline and fine-tuned DistilBERT")
+predict_parser.add_argument("-f", "--file", required=True, help="Path to trained model folder")
+predict_parser.add_argument("--title", required=True, help="Job title to classify")
+predict_parser.add_argument("--description", required=True, help="Job description to classify")
+
+# Parse the arguments
+args = parser.parse_args()
+
+if args.subcommand == "predict":
+    print(f"""===\n\nPredicting with {args.model} using
+          title '{args.title[:50]}{'...' if len(args.title) > 50 else ''}' and
+          description '{args.description[:50]}{'...' if len(args.description) > 50 else ''}'""")
+
+    if args.model == ModelEnum.SVM:
+        print("Loading SVM model...")
+        from fake_job_detector.models import BaselineSVMModel
+        model = BaselineSVMModel()
+        model.load_model(args.file)
+
+    elif args.model == ModelEnum.DISTILBERT:
+        print("Loading DistilBERT model for CPU inference...")
+        from fake_job_detector.models import DistilBERTBaseModel
+        model = DistilBERTBaseModel(pretrained_model=args.file, cpu=True)
+
+    print(f"===\n\nJob is {'fake' if model(args.title, args.description) else 'real'}")

fake_job_detector/dataset.py

@@ -0,0 +1,227 @@
+"""
+Module for storing the Dataset class which will compartmentalize things like the
+train-test split and shuffling (if needed.)
+
+Feel free to extend the class if you want to implement something specific to
+your method like dataset shuffling and batching for DL methods.
+"""
+
+from typing import Tuple, Optional
+
+from scipy.sparse import hstack
+import pandas as pd
+from datasets import (
+    load_dataset,
+    Dataset,
+    DatasetDict,
+    ClassLabel,
+    Features,
+    Value
+)
+from sklearn.model_selection import train_test_split
+from sklearn.feature_extraction.text import TfidfVectorizer
+
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, DataCollatorWithPadding
+
+from .utils import NegClassRandomSampler
+
+
+class JobDataset:
+    """
+    Wrapper around the AEGEAN dataset
+    """
+    def __init__(self,
+                 batch_size: int = 16,
+                 train_test_split: Tuple[float, float, float] = (0.7, 0.1, 0.2)):
+        _dataset = load_dataset("victor/real-or-fake-fake-jobposting-prediction")
+        self._dataset: pd.DataFrame = _dataset['train'].to_pandas()
+        self._dataset['fraudulent'] = self._dataset['fraudulent'].astype(int)
+        self._size: int = len(self._dataset)
+        self._batch_size = batch_size
+        self.clean_dataset()
+        self.add_features()
+        self.set_train_test_split(*train_test_split)
+
+    def clean_dataset(self):
+        """
+        Clean up the dataset.
+        - Fills None strings
+        - Converts label to an int
+        """
+        self._dataset[["title", "description"]] = \
+            self._dataset[["title", "description"]].fillna("")
+
+    def add_features(self):
+        """
+        Computes additional features on its own data.
+        - Concatenates the strings
+        """
+        self._dataset["full_text"] = \
+            self._dataset[["title", "description"]].agg(' '.join, axis=1)
+
+    def set_train_test_split(self,
+                             train_fr: float,
+                             eval_fr: float,
+                             test_fr: float,
+                             seed: int = 42):
+        """
+        Sets the train-test split. A seed is used for consistency.
+        """
+        eval_fr = eval_fr / (train_fr + eval_fr)
+        _train_df, self._test_df = \
+            train_test_split(self._dataset, test_size=test_fr, random_state=seed)
+        self._train_df, self._eval_df = \
+            train_test_split(_train_df, test_size=eval_fr, random_state=seed)
+
+    # Functions for getting the training, eval, and test dataset
+    # The format of the dataset will depend on the model, so I'll leave this unimplemented for now
+        
+    def get_training_set(self):
+        raise NotImplementedError
+
+    def get_validation_set(self):
+        raise NotImplementedError
+
+    def get_test_set(self):
+        raise NotImplementedError
+
+
+class SVMJobDataset(JobDataset):
+    def __init__(self, vectorizer_params: Optional[dict] = None):
+        super().__init__()
+        if vectorizer_params is None:
+            vectorizer_params = {
+                'lowercase': True,
+                'stop_words': 'english',
+                'max_features': 1_000
+            }
+        self._title_vectorizer = TfidfVectorizer(**vectorizer_params)
+        self._description_vectorizer = TfidfVectorizer(**vectorizer_params)
+
+    def vectorize(self):
+        self._train_set = hstack([
+            self._title_vectorizer.fit_transform(self._train_df["title"]),
+            self._description_vectorizer.fit_transform(self._train_df["description"])
+        ])
+        self._eval_set = hstack([
+            self._title_vectorizer.transform(self._eval_df["title"]),
+            self._description_vectorizer.transform(self._eval_df["description"])
+        ])
+        self._test_set = hstack([
+            self._title_vectorizer.transform(self._test_df["title"]),
+            self._description_vectorizer.transform(self._test_df["description"])
+        ])
+
+    def get_training_set(self):
+        return self._train_set, self._train_df["fraudulent"]
+
+    def get_validation_set(self):
+        return self._eval_set, self._eval_df["fraudulent"]
+
+    def get_test_set(self):
+        return self._test_set, self._test_df["fraudulent"]
+
+
+class HuggingFaceJobDataset(JobDataset):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._hf_dataset = None
+        self._tokenized_dataset = None
+        self._tokenizer = None
+        self._data_collator = None
+
+        self.set_hf_dataset_dict()
+        self.set_tokenized_hf_dataset()
+
+    def set_hf_dataset_dict(self, recompute: bool = False):
+        if (self._hf_dataset is not None) and (not recompute):
+            print("HF dataset already exists, recompute not set to True, returning")
+            return
+
+        hf_dataset = DatasetDict()
+
+        # Set the splits
+        features = Features({
+            "full_text": Value("string"),
+            "fraudulent": ClassLabel(num_classes=2, names=[0,1]),
+            "__index_level_0__": Value("uint32")
+        })
+        columns = ["full_text", "fraudulent"]
+        hf_dataset['train'] = Dataset.from_pandas(self._train_df[columns], features=features)
+        hf_dataset['validation'] = Dataset.from_pandas(self._eval_df[columns], features=features)
+        hf_dataset['test'] = Dataset.from_pandas(self._test_df[columns], features=features)
+
+        # Set proper names
+        hf_dataset = hf_dataset \
+            .rename_column("full_text", "text") \
+            .rename_column("fraudulent", "labels")
+
+        # Remove the index
+        hf_dataset = hf_dataset.remove_columns("__index_level_0__")
+
+        self._sampler_ratio: float = None
+        self._hf_dataset = hf_dataset
+
+    def set_tokenized_hf_dataset(self, recompute: bool = False):
+        if (self._data_collator is not None) and (self._tokenized_dataset is not None) and (not recompute):
+            print("Tokenized dataset already exists, recompute not set to True, returning")
+            return
+
+        self._tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+        hf_dataset = self._hf_dataset
+
+        tokenized_dataset = DatasetDict()
+        tokenized_dataset["train"] = hf_dataset["train"].map(self._preprocess_function, batched=True)
+        tokenized_dataset["validation"] = hf_dataset["validation"].map(self._preprocess_function, batched=True)
+        tokenized_dataset["test"] = hf_dataset["test"].map(self._preprocess_function, batched=True)
+
+        self._data_collator = DataCollatorWithPadding(tokenizer=self._tokenizer)
+        self._tokenized_dataset = tokenized_dataset
+
+    def set_random_sampler_ratio(self, neg_class_ratio: float = 0.2):
+        """
+        For randomly subsampling the negative class
+        """
+        self._sampler_ratio = neg_class_ratio
+
+    def _get_set(self, dataset_name, dataloader, subsample):
+        _ds = self._tokenized_dataset[dataset_name]
+        if subsample:
+            sample_size = min(512, len(self._tokenized_dataset[dataset_name]))
+            _ds = self._tokenized_dataset[dataset_name].shuffle(seed=42).select(list(range(sample_size)))
+        if dataloader:
+            _dst = _ds.remove_columns("text")
+            if self._sampler_ratio is None:
+                _ds = DataLoader(
+                    _dst,
+                    shuffle=True,
+                    batch_size=self._batch_size,
+                    collate_fn=self._data_collator
+                )
+            else:
+                _ds = DataLoader(
+                    _dst,
+                    batch_size=self._batch_size,
+                    collate_fn=self._data_collator,
+                    sampler=NegClassRandomSampler(_dst, self._sampler_ratio)
+                )
+        return _ds
+
+    def get_training_set(self, dataloader=True, subsample=False):
+        return self._get_set("train", dataloader, subsample)
+
+    def get_validation_set(self, dataloader=True, subsample=False):
+        return self._get_set("validation", dataloader, subsample)
+
+    def get_test_set(self, dataloader=True, subsample=False):
+        return self._get_set("test", dataloader, subsample)
+
+    def get_preprocessors(self):
+        return {
+            "tokenizer": self._tokenizer,
+            "data_collator": self._data_collator
+        }
+
+    def _preprocess_function(self, examples):
+            return self._tokenizer(examples["text"], padding="max_length", truncation=True)

fake_job_detector/models.py

@@ -0,0 +1,263 @@
+"""
+Module for storing the Model class, which can be used for wrapping sklearn or
+PyTorch models. This is more so that evaluation can be abstracted.
+"""
+
+import pickle
+import os
+from abc import ABC, abstractmethod
+from typing import Optional
+
+import numpy as np
+import matplotlib.pyplot as plt
+
+# Sklearn
+from scipy.sparse import hstack
+from sklearn.svm import SVC
+
+# Torch
+import torch
+from torch.optim import AdamW
+from transformers import (
+    AutoModelForSequenceClassification,
+    AutoTokenizer,
+    get_scheduler
+)
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+
+from .dataset import JobDataset, SVMJobDataset, HuggingFaceJobDataset
+from .utils import FocalLoss, compute_metrics
+
+
+class Model(ABC):
+
+    # Saving and loading
+    @abstractmethod
+    def save_model(self, path: str, *args):
+        """Save the model into a serialized format (e.g. pickle, tensors)"""
+        pass
+
+    @abstractmethod
+    def load_model(self, path: str, *args):
+        """Loads the model from the serialized format"""
+        pass
+
+    # Training, inference, evaluation
+    @abstractmethod
+    def fit(self, dataset: JobDataset):
+        """Given the dataset class, train the underlying model"""
+        pass
+
+    @abstractmethod
+    def evaluate(self, dataset: JobDataset):
+        """Given the dataset class, output the evaluation metrics"""
+        pass
+
+    @abstractmethod
+    def __call__(self, *args, **kwargs):
+        """Given model inputs, predict the test set labels"""
+        pass
+
+
+
+class DistilBERTBaseModel(Model):
+    def __init__(self,
+                 pretrained_model="distilbert-base-uncased",
+                 num_labels=2,
+                 freeze=False,
+                 class_frequencies: Optional[torch.Tensor] = None,
+                 cpu=False):
+        self._device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        if cpu:
+            self._device = torch.device("cpu")
+        print("Torch device: ", repr(self._device))
+
+        self._model = AutoModelForSequenceClassification.from_pretrained(
+            pretrained_model, num_labels=num_labels
+        ).to(self._device)
+        self._tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")
+
+        # Initially we trained it with the earlier layers frozen to try and
+        # speed up training, however we eventually undid this once we've
+        # established the training loop.
+        if freeze:
+            self.freeze_layers()
+
+        # If class frequencies were provided, use them to construct the focal
+        # loss formulation
+        self._loss = None
+        if class_frequencies is not None:
+            print(f"Loading a-balanced focal loss with weights {str(class_frequencies)}")
+            self._loss = FocalLoss(
+                class_frequencies=class_frequencies
+            )
+
+        # Set defaults, calling this after the fact to re-set parameters is
+        # simple enough
+        self.set_training_args()
+
+    def freeze_layers(self, layer_prefixes: Optional[set] = None):
+        """
+        Freezes certain layers by prefixes in order to focus training on only
+        certain layers.
+        """
+
+        if layer_prefixes is None:
+            layer_prefixes = set([
+            "distilbert.embeddings",
+            "distilbert.transformer.layer.0",
+            "distilbert.transformer.layer.1",
+            "distilbert.transformer.layer.2",
+            "distilbert.transformer.layer.3",
+        ])
+
+        for name, params in self._model.named_parameters():
+            if any(prefix for prefix in layer_prefixes if name.startswith(prefix)):
+                params.requires_grad = False
+
+    def set_training_args(self, **training_args):
+        training_args.setdefault("output_dir", "../models/DistilBERTBase")
+        training_args.setdefault("learning_rate", 2e-5)
+        training_args.setdefault("per_device_train_batch_size", 16)
+        training_args.setdefault("per_device_eval_batch_size", 16)
+        training_args.setdefault("num_train_epochs", 3)
+        training_args.setdefault("weight_decay", 0.01)
+        training_args.setdefault("save_strategy", "epoch")
+        training_args.setdefault("evaluation_strategy", "epoch")
+        training_args.setdefault("logging_strategy", "epoch")
+
+        self._train_args = training_args
+ 
+    def save_model(self, path, checkpoint_name: str = "checkpoint"):
+        path = os.path.join(path, checkpoint_name)
+        self._model.save_pretrained(path)
+
+    def load_model(self, path):
+        self._model = AutoModelForSequenceClassification \
+            .from_pretrained(path) \
+            .to(self._device)
+
+    def fit(self,
+            dataset: HuggingFaceJobDataset,
+            subsample: bool = False,
+            plot_loss: bool = False,
+            eval_loss: bool = False):
+
+        # Set up optimizer and LR scheduler
+        train_dataloader = dataset.get_training_set(dataloader=True, subsample=subsample)
+        eval_dataloader = dataset.get_validation_set(dataloader=True, subsample=subsample)
+        num_epochs = self._train_args["num_train_epochs"]
+        num_batches = len(train_dataloader)
+        num_training_steps = num_epochs * num_batches
+        optimizer = AdamW(self._model.parameters(), lr=5e-5)
+        lr_scheduler = get_scheduler(
+            name="linear",
+            optimizer=optimizer,
+            num_warmup_steps=0,
+            num_training_steps=num_training_steps
+        )
+        progress_bar = tqdm(range(num_training_steps))
+        losses = []
+        eval_losses = []
+
+        # Training loop
+        self._model.train()
+        for epoch in range(num_epochs):
+            epoch_loss = 0.0
+            for batch in train_dataloader:
+                batch = {k: v.to(self._device) for k, v in batch.items()}
+                outputs = self._model(**batch)
+                if self._loss is None:
+                    loss = outputs.loss
+                else:
+                    logits = outputs.logits
+                    labels = batch["labels"]
+                    scores = torch.softmax(logits, dim=-1)[:len(labels), 1]
+                    loss = self._loss(scores, labels)
+                loss.backward()
+
+                optimizer.step()
+                lr_scheduler.step()
+                optimizer.zero_grad()
+                progress_bar.update(1)
+
+                epoch_loss += loss.item()
+                losses.append(loss.item())
+            avg_loss = epoch_loss / num_batches
+            print(f"Epoch {epoch+1} avg_loss: {avg_loss:.5f}")
+
+            if eval_loss:
+                eval_epoch_loss = 0.0
+                num_eval_batches = len(eval_dataloader)
+                for batch in eval_dataloader:
+                    batch = {k: v.to(self._device) for k, v in batch.items()}
+                    with torch.no_grad():
+                        outputs = self._model(**batch)
+                    loss = outputs.loss
+                    eval_epoch_loss += loss.item()
+                    eval_losses.append(loss.item())
+                avg_loss = eval_epoch_loss / num_eval_batches
+                print(f"        eval avg_loss: {avg_loss:.5f}")
+
+        # Plot the loss if requested
+        # Note that this is a moving average of the per-batch loss, which is
+        # different from the usual per-epoch loss, as we only fine-tune for a
+        # small number of epochs
+        if plot_loss:
+            kernel = np.ones(8) / 8
+            losses = np.convolve(np.array(losses), kernel, mode='valid')
+            fig, ax = plt.subplots(figsize=(10, 5))
+            ax.plot(losses, label='Training Loss (MA-8)')
+            if eval_losses:
+                ax2 = ax.twiny()
+                eval_losses = np.convolve(np.array(eval_losses), kernel, mode='valid')
+                ax2.plot(eval_losses, color='orange', label='Eval Loss (MA-8)')
+                ax2.legend()
+            ax.set_xlabel('Batch')
+            ax.set_ylabel('Average Loss')
+            ax.set_title('Loss over Batches')
+            ax.legend()
+            fig.show()
+
+    def evaluate(self, dataset: DataLoader, get_raw_results: bool = False, plot_pr_curve: bool = True):
+        self._model.eval()
+        targs_list = []
+        score_list = []
+        preds_list = []
+
+        for batch in tqdm(dataset):
+            batch = {k: v.to(self._device) for k, v in batch.items()}
+            with torch.no_grad():
+                outputs = self._model(**batch)
+
+            logits = outputs.logits
+            labels = batch["labels"]
+            scores = torch.softmax(logits, dim=-1)[:len(labels), 1]
+            predictions = torch.argmax(logits, dim=-1)
+
+            targs_list.append(labels)
+            score_list.append(scores)
+            preds_list.append(predictions)
+
+        targs = torch.concat(targs_list).cpu()
+        scores = torch.concat(score_list).cpu()
+        preds = torch.concat(preds_list).cpu()
+
+        if get_raw_results:
+            return targs, scores, preds
+        else:
+            return compute_metrics(targs, scores, preds, plot_pr_curve)
+
+    def __call__(self, title: str, description: str) -> bool:
+        inputs = self._tokenizer(
+            title + " " + description,
+            return_tensors="pt",
+            truncation=True,
+            padding=True
+        ).to(self._device)
+
+        with torch.inference_mode():
+            outputs = self._model(**inputs)
+        predictions = torch.argmax(outputs.logits, dim=-1).tolist()[0]
+        return bool(predictions)

fake_job_detector/train.py

@@ -0,0 +1,56 @@
+import argparse
+from pathlib import Path
+
+from .dataset import JobDataset, HuggingFaceJobDataset
+from .utils import compute_metrics
+from .models import Model, DistilBERTBaseModel
+
+
+def train_model_from_cli(args):
+    model_name = args.model
+    model_dir = args.model_dir
+    experiment = args.experiment
+
+    if model_name == "distilbert":
+        model_title = "DistilBERTBase"
+        model_path = Path(model_dir, model_title, experiment)
+        model_path.mkdir(parents=True, exist_ok=True)
+        subsample = args.subsample
+
+        training_args = {
+            "learning_rate": args.learning_rate,
+            "per_device_train_batch_size": args.batch_size,
+            "per_device_eval_batch_size": args.batch_size,
+            "num_train_epochs": args.epochs,
+            "weight_decay": args.weight_decay,
+            "save_steps": args.save_steps,
+        }
+
+        dataset = HuggingFaceJobDataset()
+
+        model = DistilBERTBaseModel()
+        model.set_training_args(**training_args)
+        model.fit(dataset, subsample=subsample)
+        print(model.evaluate(subsample=subsample))
+
+
+def main():
+    parser = argparse.ArgumentParser(description='Trains the fake job detector model.')
+    parser.add_argument("model", type=str, choices=["distilbert"], help="Which model to train.")
+    parser.add_argument("--model_dir", type=str, default="./models", help="Where to store the models after training.")
+    parser.add_argument("--experiment", type=str, default="base", help="Name of experiment.")
+
+    distilbert_group = parser.add_argument_group("DistilBERT training arguments")
+    distilbert_group.add_argument("--learning_rate", type=float, default=2e-5, help="Learning rate of model.")
+    distilbert_group.add_argument("--batch_size", type=int, default=16, help="Batch size when training or evaluating the model.")
+    distilbert_group.add_argument("--epochs", type=int, default=3, help="Number of epochs to train the model.")
+    distilbert_group.add_argument("--weight_decay", type=float, default=0.01, help="Weight decay induced.")
+    distilbert_group.add_argument("--save_steps", type=int, default=5, help="Number of training steps in between checkpoints.")
+    distilbert_group.add_argument("--subsample", type=bool, default=False, help="Whether or not to use only a subsample.")
+
+    args = parser.parse_args()
+    train_model_from_cli(args)
+
+
+if __name__ == "__main__":
+    main()

fake_job_detector/utils.py

@@ -0,0 +1,104 @@
+"""
+Module for defining utilities for training such as the negative class sampler
+and focal loss function.
+"""
+
+import numpy as np
+from sklearn.metrics import (
+    precision_recall_fscore_support,
+    precision_recall_curve,
+    auc,
+    PrecisionRecallDisplay
+)
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import Sampler
+
+
+def compute_metrics(targs, scores, preds, plot_pr_curve: bool = True):
+    precision, recall, f1, _ = precision_recall_fscore_support(targs, preds, average="binary")
+    prs, rcs, _ = precision_recall_curve(targs, scores)
+
+    if plot_pr_curve:
+        display = PrecisionRecallDisplay.from_predictions(
+            targs, scores, plot_chance_level=True
+        )
+        display.ax_.set_title("Precision-Recall curve of subsample")
+        display.figure_.show()
+
+    try:
+        pr_auc = auc(prs, rcs)
+    except ValueError:
+        print("Warning: curve is non-monotonic, returning None")
+        pr_auc = None
+
+    return {
+        'precision': precision,
+        'recall': recall,
+        'f1': f1,
+        'pr_auc': pr_auc
+    }
+
+
+class FocalLoss(nn.Module):
+    def __init__(self, class_frequencies: torch.Tensor, gamma: int = 2):
+        super(FocalLoss, self).__init__()
+        self.alpha = (1 / class_frequencies).to(
+            torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        )
+        self.alpha = (self.alpha / self.alpha.sum())
+        self.gamma = gamma
+
+    def forward(self, inputs, targets):
+        alpha_targets = self.alpha[targets]
+        if inputs.data.type() != targets.data.type():
+            targets = targets.type_as(inputs.data)
+        if self.alpha.type() != inputs.data.type():
+            self.alpha = self.alpha.type_as(inputs.data)
+        ce_loss = F.cross_entropy(inputs, targets, reduction='none')
+        pt = torch.exp(-ce_loss)
+        loss = (alpha_targets * (1 - pt) ** self.gamma * ce_loss).mean()
+        return loss
+
+
+class NegClassRandomSampler(Sampler):
+    """
+    Dataloader Sampler that subsamples the negative class after each epoch.
+    The idea is that we want to keep the positive samples but select a random
+    subset of negative samples each epoch for a fresh set.
+
+    With the current settings, the sampling is done without replacement, and we
+    end up with a roughly 20% data imbalance, which should hopefully be more
+    manageable.
+    """
+
+    def __init__(self, data_source, neg_class_ratio: float = 0.2, seed: int = 42):
+        self._random_gen = np.random.default_rng(seed)
+        self.data_source = data_source
+        self._neg_class_ratio = neg_class_ratio
+
+        # Get indices of the positive and negative cases
+        self._pos_indices = np.argwhere(np.array(data_source['labels']) == 1).flatten()
+        self._neg_indices = np.argwhere(np.array(data_source['labels']) == 0).flatten()
+        self._neg_num_samples = int(len(self._neg_indices) * neg_class_ratio)
+        self._pos_num_samples = len(self._pos_indices)
+
+    @property
+    def num_samples(self):
+        return self._pos_num_samples + self._neg_num_samples
+
+    def __iter__(self):
+        """
+        Each time an iteration of this is requested, the resampling is done.
+        """
+        _neg_samples = self._random_gen.choice(self._neg_indices, self._neg_num_samples, replace=False)
+        _samples = np.concatenate((_neg_samples, self._pos_indices), axis=0)
+        self._random_gen.shuffle(_samples)
+        if (len(_samples) != len(self)):
+            raise ValueError("Length of output samples (%d) does not match expected (%d)", len(_samples), len(self))
+        return iter(_samples.tolist())
+
+    def __len__(self):
+        return self.num_samples
+