Upload Wav2Vec2ForMultiHeadMultiLabelClassification

modeling_wav2vec2multihead.py

@@ -8,6 +8,7 @@ from transformers import (
     Wav2Vec2PreTrainedModel,
 )
 from transformers.utils import ModelOutput
+from .configuration_wav2vec2multihead import Wav2Vec2MultiHeadConfig
 
 
 @dataclass
@@ -21,6 +22,14 @@ class Wav2Vec2MultiHeadMultiLabelOutput(ModelOutput):
 
 
 class Wav2Vec2ForMultiHeadMultiLabelClassification(Wav2Vec2PreTrainedModel):
+    
+    """Wav2Vec2ForMultiHeadMultiLabelClassification is a model for multi-label classification using Wav2Vec2 using multiple classifier heads. Three classifier heads are hard-coded for three different tasks, such as action, object, and location classification in FSC-IC dataset.
+
+    Returns:
+        Wav2Vec2MultiHeadMultiLabelOutput: Contains the loss and logits for each of the three tasks, as well as hidden states and attentions if requested.
+    """
+
+    config_class = Wav2Vec2MultiHeadConfig
     def __init__(self, config):
         super().__init__(config)
         self.wav2vec2 = Wav2Vec2Model(config)