batch / output objects
@@ -17,7 +17,7 @@ See below for documentation on each specific tool.
"""
from torch import optim
-from collections import deque, namedtuple
+from collections import deque
import json
import random
import numpy as np
@@ -25,7 +25,7 @@ from parlai.core.agents import Agent
from parlai.core.build_data import modelzoo_path
from parlai.core.dict import DictionaryAgent
from parlai.core.utils import (
- set_namedtuple_defaults, argsort, padded_tensor, warn_once, round_sigfigs
+ AttrDict, argsort, padded_tensor, warn_once, round_sigfigs
)
from parlai.core.distributed_utils import is_primary_worker
@@ -35,13 +35,8 @@ except ImportError:
raise ImportError('Need to install Pytorch: go to pytorch.org')
-Batch = namedtuple('Batch', [
- 'text_vec', 'text_lengths', 'label_vec', 'label_lengths', 'labels',
- 'valid_indices', 'candidates', 'candidate_vecs', 'image',
- 'memory_vecs', 'observations'
-])
-set_namedtuple_defaults(Batch, default=None)
-Batch.__doc__ = """
+class Batch(AttrDict):
+ """
Batch is a namedtuple containing data being sent to an agent.
This is the input type of the train_step and eval_step functions.
@@ -104,10 +99,23 @@ parent function to set up these fields as a base.
the original observations in the batched order
"""
-
-Output = namedtuple('Output', ['text', 'text_candidates'])
-set_namedtuple_defaults(Output, default=None)
-Output.__doc__ = """
+ def __init__(self, text_vec=None, text_lengths=None,
+ label_vec=None, label_lengths=None, labels=None,
+ valid_indices=None,
+ candidates=None, candidate_vecs=None,
+ image=None, memory_vecs=None, observations=None,
+ **kwargs):
+ super().__init__(
+ text_vec=text_vec, text_lengths=text_lengths,
+ label_vec=label_vec, label_lengths=label_lengths, labels=labels,
+ valid_indices=valid_indices,
+ candidates=candidates, candidate_vecs=candidate_vecs,
+ image=image, memory_vecs=memory_vecs, observations=observations,
+ **kwargs)
+
+
+class Output(AttrDict):
+ """
Output is a namedtuple containing agent predictions.
This is the expected return type of the train_step and eval_step functions,
@@ -123,6 +131,9 @@ though agents can choose to return None if they do not want to answer.
each sub-list is an ordered ranking of strings, of variable length.
"""
+ def __init__(self, text=None, text_candidates=None, **kwargs):
+ super().__init__(text=text, text_candidates=text_candidates, **kwargs)
+
class TorchAgent(Agent):
"""A provided base agent for any model that wants to use Torch.
