Update handler.py

handler.py

@@ -1,10 +1,12 @@
 from transformers import AutoConfig, Wav2Vec2Processor
+from transformers.file_utils import ModelOutput
+from dataclasses import dataclass
 from torch import nn
 import torch
 import io
 import torchaudio
 import torch.nn.functional as F
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Optional, Tuple
 from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
 import requests
 import tempfile
@@ -16,6 +18,18 @@ from transformers.models.wav2vec2.modeling_wav2vec2 import (
 )
 
 
+
+
+
+
+@dataclass
+class SpeechClassifierOutput(ModelOutput):
+    loss: Optional[torch.FloatTensor] = None
+    logits: torch.FloatTensor = None
+    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
+    attentions: Optional[Tuple[torch.FloatTensor]] = None
+
+
 class Wav2Vec2ClassificationHead(nn.Module):
     """Head for wav2vec classification task."""