Include hidden weights in the ClassificationOutput Tuple

[srijiths]

Hi @VuBacktracking Thank you so much for this implementation. It works well. Now, i like to include the hidden_states in the ClassificationOutput named tuple. When i do this, during evaluation, its getting in to tensor shape issues. Can you please tell me how to add hidden_states also in ClassificationOutput named tuple.

I am getting below issue, when i am adding hidden_states in ClassificationOutput tuple in MambaTextClassification class

---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
<ipython-input-129-3435b262f1ae> in <cell line: 1>()
----> 1 trainer.train()

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in train(self, resume_from_checkpoint, trial, ignore_keys_for_eval, **kwargs)
   1936                 hf_hub_utils.enable_progress_bars()
   1937         else:
-> 1938             return inner_training_loop(
   1939                 args=args,
   1940                 resume_from_checkpoint=resume_from_checkpoint,

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in _inner_training_loop(self, batch_size, args, resume_from_checkpoint, trial, ignore_keys_for_eval)
   2354                     self.control = self.callback_handler.on_step_end(args, self.state, self.control)
   2355 
-> 2356                     self._maybe_log_save_evaluate(tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2357                 else:
   2358                     self.control = self.callback_handler.on_substep_end(args, self.state, self.control)

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in _maybe_log_save_evaluate(self, tr_loss, grad_norm, model, trial, epoch, ignore_keys_for_eval)
   2802         metrics = None
   2803         if self.control.should_evaluate:
-> 2804             metrics = self._evaluate(trial, ignore_keys_for_eval)
   2805 
   2806         if self.control.should_save:

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler)
   2759 
   2760     def _evaluate(self, trial, ignore_keys_for_eval, skip_scheduler=False):
-> 2761         metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
   2762         self._report_to_hp_search(trial, self.state.global_step, metrics)
   2763 

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in evaluate(self, eval_dataset, ignore_keys, metric_key_prefix)
   3664 
   3665         eval_loop = self.prediction_loop if self.args.use_legacy_prediction_loop else self.evaluation_loop
-> 3666         output = eval_loop(
   3667             eval_dataloader,
   3668             description="Evaluation",

/usr/local/lib/python3.10/dist-packages/transformers/trainer.py in evaluation_loop(self, dataloader, description, prediction_loss_only, ignore_keys, metric_key_prefix)
   3954                 )
   3955             else:
-> 3956                 metrics = self.compute_metrics(EvalPrediction(predictions=all_preds, label_ids=all_labels))
   3957         elif metrics is None:
   3958             metrics = {}

<ipython-input-124-be10a8cea409> in compute_metrics(eval_pred)
      5     predictions, labels = eval_pred
      6 
----> 7     predictions = np.argmax(predictions, axis=1)
      8 
      9     return eval_metrics.compute(predictions=predictions, references=labels)

/usr/local/lib/python3.10/dist-packages/numpy/core/fromnumeric.py in argmax(a, axis, out, keepdims)
   1227     """
   1228     kwds = {'keepdims': keepdims} if keepdims is not np._NoValue else {}
-> 1229     return _wrapfunc(a, 'argmax', axis=axis, out=out, **kwds)
   1230 
   1231 

/usr/local/lib/python3.10/dist-packages/numpy/core/fromnumeric.py in _wrapfunc(obj, method, *args, **kwds)
     54     bound = getattr(obj, method, None)
     55     if bound is None:
---> 56         return _wrapit(obj, method, *args, **kwds)
     57 
     58     try:

/usr/local/lib/python3.10/dist-packages/numpy/core/fromnumeric.py in _wrapit(obj, method, *args, **kwds)
     43     except AttributeError:
     44         wrap = None
---> 45     result = getattr(asarray(obj), method)(*args, **kwds)
     46     if wrap:
     47         if not isinstance(result, mu.ndarray):

ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 2 dimensions. The detected shape was (2, 5645) + inhomogeneous part.
`