Update wav2vec2speechclassification.py

Files changed (1) hide show

wav2vec2speechclassification.py CHANGED Viewed

@@ -2,7 +2,7 @@ from dataclasses import dataclass
 from typing import Optional, Tuple
 import torch
 from transformers.file_utils import ModelOutput
-from transformers import Wav2Vec2Config
 @dataclass
@@ -25,8 +25,7 @@ from transformers.models.wav2vec2.modeling_wav2vec2 import (
 class Wav2Vec2ClassificationHead(nn.Module):
     """Head for wav2vec classification task."""
-    config_class = Wav2Vec2Config
-    model_type = "wav2vec2"
     def __init__(self, config):
         super().__init__()
@@ -45,8 +44,7 @@ class Wav2Vec2ClassificationHead(nn.Module):
 class Wav2Vec2ForSpeechClassification(Wav2Vec2PreTrainedModel):
-    config_class = Wav2Vec2Config
-    model_type = "wav2vec2"
     def __init__(self, config):
         super().__init__(config)

 from typing import Optional, Tuple
 import torch
 from transformers.file_utils import ModelOutput
+from .wav2vec2fsr_config import W2V2FSRConfig
 @dataclass
 class Wav2Vec2ClassificationHead(nn.Module):
     """Head for wav2vec classification task."""
+    config_class = W2V2FSRConfig
     def __init__(self, config):
         super().__init__()
 class Wav2Vec2ForSpeechClassification(Wav2Vec2PreTrainedModel):
+    config_class = W2V2FSRConfig
     def __init__(self, config):
         super().__init__(config)