Use Activations dataclass in wrapper
@@ -2,7 +2,7 @@ from collections import namedtuple
import torch
from thinc.extra.wrappers import PyTorchWrapper, torch2xp, xp2torch
-from .util import get_pytt_model
+from .util import get_pytt_model, Activations
class PyTT_Wrapper(PyTorchWrapper):
@@ -14,10 +14,9 @@ class PyTT_Wrapper(PyTorchWrapper):
self._model = self._model.from_pretrained(name)
return self
- def __init__(self, name, out_cols=("last_hidden_state", "pooler_output")):
+ def __init__(self, name):
model = get_pytt_model(name)
PyTorchWrapper.__init__(self, model)
- self.out_cols = out_cols
@property
def nO(self):
@@ -33,21 +32,24 @@ class PyTT_Wrapper(PyTorchWrapper):
self._model.train()
y_var = self._model(ids)
self._model.training = is_training
- converted_outputs = []
- for var in y_var[: len(self.out_cols)]:
- if isinstance(var, tuple):
- converted_outputs.append(var)
- else:
- converted_outputs.append(torch2xp(var))
- output = namedtuple("pytt_outputs", self.out_cols)(*converted_outputs)
+ output = Activations.from_pytt(y_var, is_grad=False)
+ assert output.has_last_hidden_state
def backward_pytorch(dy_data, sgd=None):
y_for_bwd = []
dy_for_bwd = []
- for y, dy in zip(y_var, dy_data):
- if dy is not None:
- dy_for_bwd.append(xp2torch(dy))
- y_for_bwd.append(y)
+ if dy_data.has_last_hidden_state:
+ dy_for_bwd.append(xp2torch(dy_data.last_hidden_state))
+ y_for_bwd.append(y_var[0])
+ if dy_data.has_pooler_output:
+ dy_for_bwd.append(xp2torch(dy_data.pooler_output))
+ y_for_bwd.append(y_var[1])
+ if dy_data.has_all_hidden_states:
+ raise ValueError(
+ "Gradients on all hidden states not supported yet.")
+ if dy_data.has_all_attentions:
+ raise ValueError(
+ "Gradients on all attentions not supported yet.")
torch.autograd.backward(y_for_bwd, grad_tensors=dy_for_bwd)
if sgd is not None:
if self._optimizer is None:
@@ -56,6 +58,5 @@ class PyTT_Wrapper(PyTorchWrapper):
self._optimizer.zero_grad()
return None
- assert output.last_hidden_state is not None
self._model.eval()
return output, backward_pytorch
