diff --git a/download_datasets.sh b/download_datasets.sh
index d3b9242e..8bb2d434 100644
--- a/download_datasets.sh
+++ b/download_datasets.sh
@@ -109,6 +109,13 @@ tar -xzvf spo_GO.tar.gz
rm spo_GO.tar.gz
cd ../..
+#TREC
+mkdir -p data/trec
+cd data/trec
+gdown "https://drive.google.com/uc?id=1uptxwHDL8suhVXt85Q6-7qu5ORNYoHXb"
+tar -xzvf trec.tar.gz
+rm trec.tar.gz
+cd ../..
diff --git a/model_configs/singlelabel_classification/datasets.jsonnet b/model_configs/singlelabel_classification/datasets.jsonnet
new file mode 100644
index 00000000..cbd9ce99
--- /dev/null
+++ b/model_configs/singlelabel_classification/datasets.jsonnet
@@ -0,0 +1,13 @@
+{
+ // text datasets
+ trec: {
+ dir_name: 'trec',
+ num_labels: {
+ coarse: 6,
+ fine: 50
+ },
+ train_file: 'train_val_fold_@(0|1|2|3|4|5).jsonl',
+ validation_file: 'train_val_fold_@(6|7|8|9).jsonl',
+ test_file: 'test.jsonl',
+ },
+}
diff --git a/model_configs/singlelabel_classification/trec_bert_coarse_tasknn.jsonnet b/model_configs/singlelabel_classification/trec_bert_coarse_tasknn.jsonnet
new file mode 100644
index 00000000..657cffa2
--- /dev/null
+++ b/model_configs/singlelabel_classification/trec_bert_coarse_tasknn.jsonnet
@@ -0,0 +1,160 @@
+// Simple tasknn only model for blurb genre collection
+local test = std.extVar('TEST'); // a test run with small dataset
+local data_dir = std.extVar('DATA_DIR');
+local cuda_device = std.extVar('CUDA_DEVICE');
+local use_wandb = (if test == '1' then false else true);
+
+local dataset_name = 'trec';
+local label_granularity = 'coarse';
+local dataset_metadata = (import './datasets.jsonnet')[dataset_name];
+local num_labels = dataset_metadata.num_labels[label_granularity];
+local num_input_features = dataset_metadata.input_features;
+
+// model variables
+local ff_hidden = std.parseJson(std.extVar('ff_hidden'));
+local label_space_dim = ff_hidden;
+local task_nn_dropout = std.parseJson(std.extVar('task_nn_dropout_10x')) / 10.0;
+local ff_activation = 'softplus';
+local ff_linear_layers = std.parseJson(std.extVar('ff_linear_layers'));
+local task_nn_weight_decay = std.parseJson(std.extVar('task_nn_weight_decay'));
+local gain = (if ff_activation == 'tanh' then 5 / 3 else 1);
+local transformer_model = 'bert-base-uncased'; // huggingface name of the model
+local transformer_dim = 768;
+{
+ [if use_wandb then 'type']: 'train_test_log_to_wandb',
+ evaluate_on_test: true,
+ // Data
+ dataset_reader: {
+ type: dataset_name,
+ granularity: label_granularity,
+ [if test == '1' then 'max_instances']: 100,
+ token_indexers: {
+ x: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ },
+ },
+ tokenizer: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ max_length: 512,
+ },
+ },
+ train_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.train_file),
+ validation_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.validation_file),
+ test_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.test_file),
+
+ // Model
+ model: {
+ type: 'single-label-classification-with-infnet',
+ sampler: {
+ type: 'appending-container',
+ log_key: 'sampler',
+ constituent_samplers: [],
+ },
+ task_nn: {
+ type: 'single-label-text-classification',
+ feature_network: {
+ text_field_embedder: {
+ token_embedders: {
+ x: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ },
+ },
+ },
+ seq2vec_encoder: {
+ type: 'bert_pooler',
+ pretrained_model: transformer_model,
+ },
+ final_dropout: task_nn_dropout,
+ feedforward: {
+ input_dim: transformer_dim,
+ num_layers: ff_linear_layers,
+ activations: ([ff_activation for i in std.range(0, ff_linear_layers - 2)] + [ff_activation]),
+ hidden_dims: ff_hidden,
+ dropout: ([task_nn_dropout for i in std.range(0, ff_linear_layers - 2)] + [0]),
+ },
+ },
+ label_embeddings: {
+ embedding_dim: ff_hidden,
+ vocab_namespace: 'labels',
+ },
+ },
+ inference_module: {
+ type: 'single-label-basic',
+ log_key: 'inference_module',
+ loss_fn: {
+ type: 'single-label-ce',
+ reduction: 'mean',
+ log_key: 'cross_entropy',
+ },
+ },
+ oracle_value_function: null, // { type: 'per-instance-f1', differentiable: false }
+ score_nn: null, // no score nn for basic tasknn model
+ loss_fn: { type: 'zero' },
+ initializer: {
+ regexes: [
+ [@'.*feedforward._linear_layers.*weight', (if std.member(['tanh', 'sigmoid'], ff_activation) then { type: 'xavier_uniform', gain: gain } else { type: 'kaiming_uniform', nonlinearity: 'relu' })],
+ [@'.*feedforward._linear_layers.*bias', { type: 'zero' }],
+ ],
+ },
+ },
+ data_loader: {
+ "batch_sampler": {
+ "type": "bucket",
+ "batch_size" : 2 // effective batch size = batch_size*num_gradient_accumulation_steps
+ },
+ },
+ trainer: {
+ type: 'gradient_descent_minimax',
+ num_epochs: if test == '1' then 10 else 300,
+ grad_norm: { task_nn: 1.0 },
+ num_gradient_accumulation_steps: 16, // effective batch size = batch_size*num_gradient_accumulation_steps
+ patience: 5,
+ validation_metric: '+accuracy',
+ cuda_device: std.parseInt(cuda_device),
+ learning_rate_schedulers: {
+ task_nn: {
+ type: 'reduce_on_plateau',
+ factor: 0.5,
+ mode: 'max',
+ patience: 2,
+ verbose: true,
+ },
+ },
+ optimizer: {
+ optimizers: { // have only tasknn optmizer
+ task_nn:
+ {
+ lr: 1e-5,
+ weight_decay: task_nn_weight_decay,
+ type: 'huggingface_adamw',
+ },
+ },
+ },
+ checkpointer: {
+ keep_most_recent_by_count: 1,
+ },
+ callbacks: [
+ 'track_epoch_callback',
+ 'slurm',
+ ] + (
+ if use_wandb then [
+ {
+ type: 'wandb_allennlp',
+ sub_callbacks: [{ type: 'log_best_validation_metrics', priority: 100 }],
+ save_model_archive: false,
+ // watch_model: false,
+ should_log_parameter_statistics: false,
+ },
+ ]
+ else []
+ ),
+ inner_mode: 'score_nn',
+ num_steps: { task_nn: 1, score_nn: 1 },
+ },
+}
diff --git a/model_configs/singlelabel_classification/trec_bert_fine_tasknn.jsonnet b/model_configs/singlelabel_classification/trec_bert_fine_tasknn.jsonnet
new file mode 100644
index 00000000..ff24e119
--- /dev/null
+++ b/model_configs/singlelabel_classification/trec_bert_fine_tasknn.jsonnet
@@ -0,0 +1,160 @@
+// Simple tasknn only model for blurb genre collection
+local test = std.extVar('TEST'); // a test run with small dataset
+local data_dir = std.extVar('DATA_DIR');
+local cuda_device = std.extVar('CUDA_DEVICE');
+local use_wandb = (if test == '1' then false else true);
+
+local dataset_name = 'trec';
+local label_granularity = 'fine';
+local dataset_metadata = (import './datasets.jsonnet')[dataset_name];
+local num_labels = dataset_metadata.num_labels[label_granularity];
+local num_input_features = dataset_metadata.input_features;
+
+// model variables
+local ff_hidden = std.parseJson(std.extVar('ff_hidden'));
+local label_space_dim = ff_hidden;
+local task_nn_dropout = std.parseJson(std.extVar('task_nn_dropout_10x')) / 10.0;
+local ff_activation = 'softplus';
+local ff_linear_layers = std.parseJson(std.extVar('ff_linear_layers'));
+local task_nn_weight_decay = std.parseJson(std.extVar('task_nn_weight_decay'));
+local gain = (if ff_activation == 'tanh' then 5 / 3 else 1);
+local transformer_model = 'bert-base-uncased'; // huggingface name of the model
+local transformer_dim = 768;
+{
+ [if use_wandb then 'type']: 'train_test_log_to_wandb',
+ evaluate_on_test: true,
+ // Data
+ dataset_reader: {
+ type: dataset_name,
+ granularity: label_granularity,
+ [if test == '1' then 'max_instances']: 100,
+ token_indexers: {
+ x: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ },
+ },
+ tokenizer: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ max_length: 512,
+ },
+ },
+ train_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.train_file),
+ validation_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.validation_file),
+ test_data_path: (data_dir + '/' + dataset_metadata.dir_name + '/' +
+ dataset_metadata.test_file),
+
+ // Model
+ model: {
+ type: 'single-label-classification-with-infnet',
+ sampler: {
+ type: 'appending-container',
+ log_key: 'sampler',
+ constituent_samplers: [],
+ },
+ task_nn: {
+ type: 'single-label-text-classification',
+ feature_network: {
+ text_field_embedder: {
+ token_embedders: {
+ x: {
+ type: 'pretrained_transformer',
+ model_name: transformer_model,
+ },
+ },
+ },
+ seq2vec_encoder: {
+ type: 'bert_pooler',
+ pretrained_model: transformer_model,
+ },
+ final_dropout: task_nn_dropout,
+ feedforward: {
+ input_dim: transformer_dim,
+ num_layers: ff_linear_layers,
+ activations: ([ff_activation for i in std.range(0, ff_linear_layers - 2)] + [ff_activation]),
+ hidden_dims: ff_hidden,
+ dropout: ([task_nn_dropout for i in std.range(0, ff_linear_layers - 2)] + [0]),
+ },
+ },
+ label_embeddings: {
+ embedding_dim: ff_hidden,
+ vocab_namespace: 'labels',
+ },
+ },
+ inference_module: {
+ type: 'single-label-basic',
+ log_key: 'inference_module',
+ loss_fn: {
+ type: 'single-label-ce',
+ reduction: 'mean',
+ log_key: 'cross_entropy',
+ },
+ },
+ oracle_value_function: null, // { type: 'per-instance-f1', differentiable: false }
+ score_nn: null, // no score nn for basic tasknn model
+ loss_fn: { type: 'zero' },
+ initializer: {
+ regexes: [
+ [@'.*feedforward._linear_layers.*weight', (if std.member(['tanh', 'sigmoid'], ff_activation) then { type: 'xavier_uniform', gain: gain } else { type: 'kaiming_uniform', nonlinearity: 'relu' })],
+ [@'.*feedforward._linear_layers.*bias', { type: 'zero' }],
+ ],
+ },
+ },
+ data_loader: {
+ "batch_sampler": {
+ "type": "bucket",
+ "batch_size" : 2 // effective batch size = batch_size*num_gradient_accumulation_steps
+ },
+ },
+ trainer: {
+ type: 'gradient_descent_minimax',
+ num_epochs: if test == '1' then 10 else 300,
+ grad_norm: { task_nn: 1.0 },
+ num_gradient_accumulation_steps: 16, // effective batch size = batch_size*num_gradient_accumulation_steps
+ patience: 5,
+ validation_metric: '+accuracy',
+ cuda_device: std.parseInt(cuda_device),
+ learning_rate_schedulers: {
+ task_nn: {
+ type: 'reduce_on_plateau',
+ factor: 0.5,
+ mode: 'max',
+ patience: 2,
+ verbose: true,
+ },
+ },
+ optimizer: {
+ optimizers: { // have only tasknn optmizer
+ task_nn:
+ {
+ lr: 1e-5,
+ weight_decay: task_nn_weight_decay,
+ type: 'huggingface_adamw',
+ },
+ },
+ },
+ checkpointer: {
+ keep_most_recent_by_count: 1,
+ },
+ callbacks: [
+ 'track_epoch_callback',
+ 'slurm',
+ ] + (
+ if use_wandb then [
+ {
+ type: 'wandb_allennlp',
+ sub_callbacks: [{ type: 'log_best_validation_metrics', priority: 100 }],
+ save_model_archive: false,
+ // watch_model: false,
+ should_log_parameter_statistics: false,
+ },
+ ]
+ else []
+ ),
+ inner_mode: 'score_nn',
+ num_steps: { task_nn: 1, score_nn: 1 },
+ },
+}
diff --git a/structured_prediction_baselines/dataset_readers/singlelabel_classification/__init__.py b/structured_prediction_baselines/dataset_readers/singlelabel_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/structured_prediction_baselines/dataset_readers/singlelabel_classification/trec.py b/structured_prediction_baselines/dataset_readers/singlelabel_classification/trec.py
new file mode 100644
index 00000000..caeabd80
--- /dev/null
+++ b/structured_prediction_baselines/dataset_readers/singlelabel_classification/trec.py
@@ -0,0 +1,184 @@
+from typing import (
+ Dict,
+ List,
+ Union,
+ Any,
+ Iterator,
+ Tuple,
+ cast,
+ Optional,
+ Iterable,
+ Literal,
+)
+import sys
+import itertools
+from wcmatch import glob
+
+if sys.version_info >= (3, 8):
+ from typing import (
+ TypedDict,
+ ) # pylint: disable=no-name-in-module
+else:
+ from typing_extensions import TypedDict, Literal, overload
+
+import logging
+import json
+from allennlp.data.dataset_readers.dataset_reader import DatasetReader
+from allennlp.data.fields import (
+ TextField,
+ MetadataField,
+ LabelField,
+)
+from allennlp.data.instance import Instance
+from allennlp.data.token_indexers import TokenIndexer
+from allennlp.data.tokenizers import Tokenizer
+from allennlp.data.tokenizers import Token
+
+import allennlp
+
+logger = logging.getLogger(__name__)
+
+
+class InstanceFields(TypedDict):
+ """Contents which form an instance"""
+
+ x: TextField
+ label: LabelField
+
+
+@DatasetReader.register("trec")
+class TRECReader(DatasetReader):
+ """
+ Single-label classification `dataset `.
+
+ """
+
+ def __init__(
+ self,
+ tokenizer: Tokenizer,
+ token_indexers: Dict[str, TokenIndexer],
+ granularity: Literal["coarse", "fine"] = "coarse",
+ **kwargs: Any,
+ ) -> None:
+ """
+ Arguments:
+ tokenizer: The tokenizer to be used.
+ token_indexers: The token_indexers to be used--one per embedder. Will usually be only one.
+ granularity: The level of granularity to be used for the labels.
+ **kwargs: Parent class args.
+ `Reference `_
+
+ """
+ super().__init__(
+ manual_distributed_sharding=True,
+ manual_multiprocess_sharding=True,
+ **kwargs,
+ )
+ self._tokenizer = tokenizer
+ self._token_indexers = token_indexers
+ self._granularity = granularity
+
+ def example_to_fields(
+ self,
+ text: str,
+ label_coarse: str,
+ label_fine: str,
+ idx: str,
+ meta: Dict = None,
+ **kwargs: Any,
+ ) -> InstanceFields:
+ """Converts a dictionary containing an example datapoint to fields that can be used
+ to create an :class:`Instance`. If a meta dictionary is passed, then it also adds raw data in
+ the meta dict.
+
+ Args:
+ text: question
+ label_coarse: coarse-grained label
+ label_fine: fine-grained label
+ idx: unique id
+ meta: None
+ **kwargs: unused
+
+ Returns:
+ Dictionary of fields with the following entries:
+ x: contains the question text
+ label: contains the class label (fine or coarse)
+
+ """
+
+ if meta is None:
+ meta = {}
+
+ meta["text"] = text
+ meta["label"] = {
+ "coarse": label_coarse,
+ "fine": label_fine
+ }
+ meta["idx"] = idx
+
+ x = TextField(
+ self._tokenizer.tokenize(text),
+ )
+ label_text = label_coarse
+ if self._granularity == "fine":
+ label_text += f"_{label_fine}" # Prepend with coarse to distinguish between multiple fine "other" labels
+ label = LabelField(label_text)
+ fields = {
+ "x": x,
+ "label": label,
+ }
+ if self._granularity == "fine":
+ fields["label_coarse"] = LabelField(label_coarse, label_namespace="coarse_labels")
+ return fields
+
+ def text_to_instance( # type:ignore
+ self,
+ text: str,
+ label_coarse: str,
+ label_fine: str,
+ idx: str,
+ **kwargs: Any,
+ ) -> Instance:
+ """Converts contents of a single raw example to :class:`Instance`.
+
+ Args:
+ text: Question text
+ label_coarse: Coarse-grained label
+ label_fine: Fine-grained label
+ idx: Identification number
+ **kwargs: unused
+
+ Returns:
+ :class:`Instance` of data
+
+ """
+ meta_dict: Dict = {}
+ main_fields = self.example_to_fields(
+ text, label_coarse, label_fine, idx, meta=meta_dict
+ )
+
+ return Instance(
+ {**cast(dict, main_fields), "meta": MetadataField(meta_dict)}
+ )
+
+ def _read(self, file_path: str) -> Iterator[Instance]:
+ """Reads a datafile to produce instances
+
+ Args:
+ file_path: Path to the data file
+
+ Yields:
+ data instances
+
+ """
+
+ for file_ in glob.glob(file_path, flags=glob.EXTGLOB):
+ with open(file_) as f:
+ for line in self.shard_iterable(f):
+ example = json.loads(line)
+ instance = self.text_to_instance(**example)
+ yield instance
+
+ def apply_token_indexers(self, instance: Instance) -> None:
+ text_field = cast(TextField, instance.fields["x"]) # no runtime effect
+ text_field._token_indexers = self._token_indexers
diff --git a/structured_prediction_baselines/metrics/__init__.py b/structured_prediction_baselines/metrics/__init__.py
index c1e90e89..d1f1da19 100644
--- a/structured_prediction_baselines/metrics/__init__.py
+++ b/structured_prediction_baselines/metrics/__init__.py
@@ -19,4 +19,7 @@
)
from .multilabel_classification_rbo import (
MultilabelClassificationRankBiasedOverlap
-)
\ No newline at end of file
+)
+from .singlelabel_classification_pairwise_difference import (
+ SinglelabelClassificationPairwiseDifference
+)
diff --git a/structured_prediction_baselines/metrics/singlelabel_classification_pairwise_difference.py b/structured_prediction_baselines/metrics/singlelabel_classification_pairwise_difference.py
new file mode 100644
index 00000000..7a0ceda5
--- /dev/null
+++ b/structured_prediction_baselines/metrics/singlelabel_classification_pairwise_difference.py
@@ -0,0 +1,53 @@
+"""Implements pairwise probability differences between classes"""
+from typing import List, Tuple, Union, Dict, Any, Optional, Literal
+
+from allennlp.training.metrics import Metric, Average
+import torch
+
+
+@Metric.register("singlelabel-pairwise-difference")
+class SinglelabelClassificationPairwiseDifference(Average):
+ """Computes the across- and within-class pairwise probability differences"""
+
+ def __init__(self, clusters: list, mode: Literal["across", "within"] = "across", drop_argmax: bool = False) -> None:
+ """
+ Args:
+ clusters: List of tensors, each containing fine-label indexes belonging to the same coarse-label
+ mode: Specifies whether the metric is computed over fine-labels from different coarse-labels ("across")
+ or the same coarse-label ("within")
+ drop_argmax: Specifies whether the highest probability label should be removed from the calculations
+ """
+
+ super().__init__()
+ self.clusters = clusters
+ self._mode = mode
+ self._drop_argmax = drop_argmax
+
+ def __call__(
+ self, predictions: torch.Tensor
+ ) -> None:
+ scores = [t.cpu() for t in self.detach_tensors(predictions)][0] # shape: (B x L)
+ # Compute pairwise probability difference between each label class
+ differences = torch.abs((scores.unsqueeze(-1) - scores.unsqueeze(1))) # shape: (B x L x L)
+ max_class = torch.argmax(scores, dim=1) # shape: (B,)
+
+ if self._mode == "within":
+ for cluster in self.clusters:
+ for _rc in torch.combinations(cluster, r=2):
+ for i, _differences in enumerate(differences):
+ if self._drop_argmax and max_class[i] in _rc:
+ # Skip if the highest probability class is either the row or column of the diff. cell
+ continue
+ diff = _differences[tuple(_rc)]
+ super().__call__(diff)
+ elif self._mode == "across":
+ for _cx in range(len(self.clusters) - 1):
+ for _cy in range(_cx + 1, len(self.clusters)):
+ _x, _y = torch.meshgrid(self.clusters[_cx], self.clusters[_cy])
+ for tup in list(zip(_x.flatten(), _y.flatten())):
+ for i, _differences in enumerate(differences):
+ if self._drop_argmax and max_class[i] in tup:
+ # Skip if the highest probability class is either the row or column of the diff. cell
+ continue
+ diff = _differences[tup]
+ super().__call__(diff)
diff --git a/structured_prediction_baselines/models/singlelabel_classification.py b/structured_prediction_baselines/models/singlelabel_classification.py
new file mode 100644
index 00000000..d7b49314
--- /dev/null
+++ b/structured_prediction_baselines/models/singlelabel_classification.py
@@ -0,0 +1,128 @@
+import logging
+from typing import List, Tuple, Union, Dict, Any, Optional
+from overrides import overrides
+from collections import defaultdict
+
+import torch
+from allennlp.models import Model
+from allennlp.training.metrics import CategoricalAccuracy
+from structured_prediction_baselines.metrics import SinglelabelClassificationPairwiseDifference
+
+from .base import ScoreBasedLearningModel
+
+logger = logging.getLogger(__name__)
+
+
+@Model.register(
+ "single-label-classification-with-infnet",
+ constructor="from_partial_objects_with_shared_tasknn",
+)
+@Model.register(
+ "single-label-classification", constructor="from_partial_objects"
+)
+@ScoreBasedLearningModel.register(
+ "single-label-classification-with-infnet",
+ constructor="from_partial_objects_with_shared_tasknn",
+)
+@ScoreBasedLearningModel.register(
+ "single-label-classification", constructor="from_partial_objects"
+)
+class SinglelabelClassification(ScoreBasedLearningModel):
+ def __init__(
+ self,
+ **kwargs: Any,
+ ) -> None:
+ super().__init__(**kwargs)
+ # Metrics
+ self.accuracy = CategoricalAccuracy()
+
+ # If coarse labels are available, track label clusters to compute pairwise probability difference metrics
+ coarse_clusters = None
+ vocab = kwargs.pop("vocab")
+ coarse_namespace = "coarse_labels"
+ self.compute_pairwise = False
+ if coarse_namespace in vocab._index_to_token.keys():
+ coarse_idx2label = vocab.get_index_to_token_vocabulary(coarse_namespace)
+ coarse_labels = coarse_idx2label.values()
+ fine_idx2label = vocab.get_index_to_token_vocabulary("labels")
+ coarse_clusters = defaultdict(list)
+ sep_char = "_" # Assuming fine labels are specified as ""
+ for idx in fine_idx2label:
+ coarse_fine = fine_idx2label[idx].split(sep_char)
+ if coarse_fine[0] in coarse_labels:
+ coarse_clusters[coarse_fine[0]].append(idx)
+ coarse_clusters = list(map(lambda x: torch.tensor(x[1]), coarse_clusters.items()))
+ if coarse_clusters is not None:
+ self.compute_pairwise = True
+ self.pairwise_diff_across = SinglelabelClassificationPairwiseDifference(coarse_clusters, mode="across")
+ self.pairwise_diff_within = SinglelabelClassificationPairwiseDifference(coarse_clusters, mode="within")
+ self.pairwise_diff_across_dropmax = SinglelabelClassificationPairwiseDifference(coarse_clusters,
+ mode="across",
+ drop_argmax=True)
+ self.pairwise_diff_within_dropmax = SinglelabelClassificationPairwiseDifference(coarse_clusters,
+ mode="within",
+ drop_argmax=True)
+
+ @overrides
+ def construct_args_for_forward(self, **kwargs: Any) -> Dict:
+ _forward_args = {}
+ _forward_args["buffer"] = self.initialize_buffer(**kwargs)
+ _forward_args["x"] = kwargs.pop("x")
+ _forward_args["labels"] = kwargs.pop("label")
+
+ # Enable the use of meta dictionary:
+ # metadata = kwargs.pop("meta")
+ # if metadata is None:
+ # raise ValueError
+ # _forward_args["buffer"]["meta"] = metadata
+
+ # Enable the use of coarse label in the forward pass:
+ # label_coarse = kwargs.pop("label_coarse")
+ # if label_coarse is not None:
+ # _forward_args["buffer"]["label_coarse"] = label_coarse.cpu()
+
+ return {**_forward_args, **kwargs}
+
+ @overrides
+ def unsqueeze_labels(self, labels: torch.Tensor) -> torch.Tensor:
+ return labels.unsqueeze(1)
+
+ @overrides
+ def squeeze_y(self, y: torch.Tensor) -> torch.Tensor:
+ return y.squeeze(1)
+
+ @torch.no_grad()
+ @overrides
+ def calculate_metrics( # type: ignore
+ self,
+ x: Any,
+ labels: torch.Tensor,
+ y_hat: torch.Tensor,
+ buffer: Dict,
+ results: Dict,
+ **kwargs: Any,
+ ) -> None:
+
+ if not self.inference_module.is_normalized:
+ y_hat_n = torch.softmax(y_hat, dim=-1)
+ else:
+ y_hat_n = y_hat
+
+ self.accuracy(y_hat_n, labels)
+
+ if self.compute_pairwise:
+ self.pairwise_diff_across(y_hat_n)
+ self.pairwise_diff_within(y_hat_n)
+ self.pairwise_diff_across_dropmax(y_hat_n)
+ self.pairwise_diff_within_dropmax(y_hat_n)
+
+ def get_true_metrics(self, reset: bool = False) -> Dict[str, float]:
+ metrics = {
+ "accuracy": self.accuracy.get_metric(reset),
+ }
+ if self.compute_pairwise:
+ metrics["pairwise_across"] = self.pairwise_diff_across.get_metric(reset)
+ metrics["pairwise_within"] = self.pairwise_diff_within.get_metric(reset)
+ metrics["pairwise_across_dropmax"] = self.pairwise_diff_across_dropmax.get_metric(reset)
+ metrics["pairwise_within_dropmax"] = self.pairwise_diff_within_dropmax.get_metric(reset)
+ return metrics
diff --git a/structured_prediction_baselines/modules/loss/singlelabel_classification/__init__.py b/structured_prediction_baselines/modules/loss/singlelabel_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/structured_prediction_baselines/modules/loss/singlelabel_classification/singlelabel_classification_cross_entropy.py b/structured_prediction_baselines/modules/loss/singlelabel_classification/singlelabel_classification_cross_entropy.py
new file mode 100644
index 00000000..f710fbd9
--- /dev/null
+++ b/structured_prediction_baselines/modules/loss/singlelabel_classification/singlelabel_classification_cross_entropy.py
@@ -0,0 +1,39 @@
+from typing import List, Tuple, Union, Dict, Any, Optional
+import torch
+from structured_prediction_baselines.modules.loss import Loss
+import numpy as np
+
+
+@Loss.register("single-label-ce")
+class SinglelabelCELoss(Loss):
+ def __init__(self, **kwargs: Any):
+ super().__init__(**kwargs)
+ self.loss_fn = torch.nn.functional.cross_entropy
+ self._loss_values = []
+
+ def _forward(
+ self,
+ x: Any,
+ labels: Optional[torch.Tensor], # (batch, 1)
+ y_hat: torch.Tensor, # (batch, 1, num_labels)
+ y_hat_extra: Optional[torch.Tensor],
+ buffer: Optional[Dict] = None,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ assert labels is not None
+ loss = self.loss_fn(y_hat.squeeze(1), labels.squeeze(-1), reduce=False) # (batch,)
+
+ self._loss_values.append(float(torch.mean(loss)))
+
+ return loss
+
+ def get_metrics(self, reset: bool = False):
+ metrics = {}
+
+ if self._loss_values:
+ metrics = {"cross_entropy_loss": np.mean(self._loss_values)}
+
+ if reset:
+ self._loss_values = []
+
+ return metrics
diff --git a/structured_prediction_baselines/modules/sampler/singlelabel_classification/__init__.py b/structured_prediction_baselines/modules/sampler/singlelabel_classification/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/structured_prediction_baselines/modules/sampler/singlelabel_classification/basic.py b/structured_prediction_baselines/modules/sampler/singlelabel_classification/basic.py
new file mode 100644
index 00000000..218fbeb7
--- /dev/null
+++ b/structured_prediction_baselines/modules/sampler/singlelabel_classification/basic.py
@@ -0,0 +1,60 @@
+from typing import List, Tuple, Union, Dict, Any, Optional, overload
+from overrides import overrides
+import torch
+from allennlp.data import TextFieldTensors
+from structured_prediction_baselines.modules.sampler import (
+ Sampler,
+ BasicSampler,
+)
+from structured_prediction_baselines.modules.score_nn import ScoreNN
+from structured_prediction_baselines.modules.oracle_value_function import (
+ OracleValueFunction,
+)
+from structured_prediction_baselines.modules.loss import Loss
+from structured_prediction_baselines.modules.multilabel_classification_task_nn import (
+ MultilabelTaskNN,
+)
+
+
+@Sampler.register("single-label-basic")
+class SinglelabelClassificationBasicSampler(BasicSampler):
+ @property
+ def is_normalized(self) -> bool:
+ return True
+
+ @property
+ def different_training_and_eval(self) -> bool:
+ return False
+
+ @overload
+ def normalize(self, y: None) -> None:
+ ...
+
+ @overload
+ def normalize(self, y: torch.Tensor) -> torch.Tensor:
+ ...
+
+ def normalize(self, y: Optional[torch.Tensor]) -> Optional[torch.Tensor]: # (batch, 1, num_labels)
+ if y is not None:
+ return torch.softmax(y, dim=-1)
+ return None
+
+ @overrides
+ def forward(
+ self,
+ x: Union[torch.Tensor, TextFieldTensors],
+ labels: Optional[torch.Tensor],
+ buffer: Dict,
+ **kwargs: Any,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor]]:
+ logits = self.inference_nn(x).unsqueeze(1) # unnormalized logits (batch, 1, ...)
+ loss = None
+ if labels is not None:
+ # compute loss for logging.
+ loss = self.loss_fn(x,
+ labels.unsqueeze(1), # (batch, num_samples or 1, ...)
+ logits,
+ logits,
+ buffer,
+ )
+ return self.normalize(logits), self.normalize(logits), loss
diff --git a/structured_prediction_baselines/modules/singlelabel_classification_task_nn.py b/structured_prediction_baselines/modules/singlelabel_classification_task_nn.py
new file mode 100644
index 00000000..4280a35b
--- /dev/null
+++ b/structured_prediction_baselines/modules/singlelabel_classification_task_nn.py
@@ -0,0 +1,56 @@
+from typing import List, Tuple, Union, Dict, Any, Optional
+from .task_nn import TaskNN, CostAugmentedLayer, TextEncoder
+from allennlp.modules.feedforward import FeedForward
+from allennlp.modules.token_embedders.embedding import Embedding
+from allennlp.data import Vocabulary, TextFieldTensors
+
+import torch.nn as nn
+import torch
+import numpy as np
+
+
+@TaskNN.register("single-label-classification")
+class SinglelabelTaskNN(TaskNN):
+ def __init__(
+ self,
+ vocab: Vocabulary,
+ feature_network: Union[FeedForward, TextEncoder],
+ label_embeddings: Embedding,
+ ):
+ super().__init__() # type:ignore
+ self.feature_network: Union[FeedForward, TextEncoder] = feature_network
+ self.label_embeddings = label_embeddings
+ assert (
+ self.label_embeddings.weight.shape[1]
+ == self.feature_network.get_output_dim() # type: ignore
+ ), (
+ f"label_embeddings dim ({self.label_embeddings.weight.shape[1]}) "
+ f"and hidden_dim of feature_network"
+ f" ({self.feature_network.get_output_dim()}) do not match."
+ )
+
+ def forward(
+ self,
+ x: Union[torch.Tensor, TextFieldTensors],
+ buffer: Optional[Dict] = None,
+ **kwargs: Any,
+ ) -> torch.Tensor:
+ features = self.feature_network(x) # (batch, hidden_dim)
+ logits = torch.matmul(features, self.label_embeddings.weight.T)
+
+ return logits # unnormalized logits of shape (batch, num_labels)
+
+
+@TaskNN.register("single-label-text-classification")
+class SinglelabelTextTaskNN(SinglelabelTaskNN):
+ def __init__(
+ self,
+ vocab: Vocabulary,
+ feature_network: TextEncoder,
+ label_embeddings: Embedding,
+ ):
+ super().__init__(
+ vocab=vocab,
+ feature_network=feature_network,
+ label_embeddings=label_embeddings,
+ )
diff --git a/tests/end2end/conftest.py b/tests/end2end/conftest.py
index 5811606f..07ac2c50 100644
--- a/tests/end2end/conftest.py
+++ b/tests/end2end/conftest.py
@@ -1,6 +1,12 @@
from .common import marks
+import nltk
+try:
+ nltk.data.find('corpora/omw-1.4')
+except LookupError:
+ nltk.download('omw-1.4')
+
def pytest_configure(config) -> None: # type: ignore
for mark in marks:
config.addinivalue_line("markers", mark)