Upload distilled speech model

README.md

@@ -30,19 +30,21 @@ A Data2Vec-style bidirectional speech encoder trained via distillation from Auri
 ## Usage
 
 ```python
-from transformers import AutoModel, AutoFeatureExtractor
+from transformers import AutoModel, Wav2Vec2FeatureExtractor
 import torch
 
 # Load model and feature extractor
 model = AutoModel.from_pretrained("TuKoResearch/AuriStreamDistill_100M40PredTeacher_librispeech960", trust_remote_code=True)
-feature_extractor = AutoFeatureExtractor.from_pretrained("TuKoResearch/AuriStreamDistill_100M40PredTeacher_librispeech960", trust_remote_code=True)
+model.eval()  # Important for inference!
+feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("TuKoResearch/AuriStreamDistill_100M40PredTeacher_librispeech960")
 
 # Prepare audio (16kHz, mono)
-audio = torch.randn(16000)  # 1 second of audio
+audio = torch.randn(16000).numpy()  # 1 second of audio
 
 # Extract features
-inputs = feature_extractor(audio, return_tensors="pt", sample_rate=16000)
-outputs = model(inputs.input_values, output_hidden_states=True)
+inputs = feature_extractor(audio, return_tensors="pt", sampling_rate=16000)
+with torch.no_grad():
+    outputs = model(inputs.input_values, output_hidden_states=True)
 
 # Get representations
 last_hidden = outputs.last_hidden_state  # (1, 50, 768) for 1 second

config.json

@@ -46,8 +46,7 @@
   "model_type": "distilled_speech",
   "auto_map": {
     "AutoConfig": "configuration_distilled_speech.DistilledSpeechConfig",
-    "AutoModel": "modeling_distilled_speech.DistilledSpeechModel",
-    "AutoFeatureExtractor": "feature_extraction_distilled_speech.DistilledSpeechFeatureExtractor"
+    "AutoModel": "modeling_distilled_speech.DistilledSpeechModel"
   },
   "architectures": [
     "DistilledSpeechModel"

modeling_distilled_speech.py

@@ -407,6 +407,9 @@ class DistilledSpeechModel(PreTrainedModel):
     representations at 50Hz (20ms stride). It returns hidden states from all
     transformer layers, making it suitable for downstream probing and finetuning.
     
+    IMPORTANT: Call model.eval() before inference to disable dropout and ensure
+    correct behavior of normalization layers.
+    
     Hidden states structure (for 12-layer model, output_hidden_states=True):
         - hidden_states[0]: Feature projection output (input to transformer)
         - hidden_states[1]: Output of transformer layer 0
@@ -416,12 +419,14 @@ class DistilledSpeechModel(PreTrainedModel):
         Total: 13 hidden states (1 embedding + 12 layers)
     
     Example usage:
-        >>> from transformers import AutoModel, AutoFeatureExtractor
+        >>> from transformers import AutoModel, Wav2Vec2FeatureExtractor
         >>> model = AutoModel.from_pretrained("your-model-name", trust_remote_code=True)
-        >>> processor = AutoFeatureExtractor.from_pretrained("your-model-name", trust_remote_code=True)
-        >>> audio = torch.randn(16000)  # 1 second of audio at 16kHz
-        >>> inputs = processor(audio, return_tensors="pt", sample_rate=16000)
-        >>> outputs = model(inputs.input_values, output_hidden_states=True)
+        >>> model.eval()  # Important for inference!
+        >>> feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained("your-model-name")
+        >>> audio = torch.randn(16000).numpy()  # 1 second of audio at 16kHz
+        >>> inputs = feature_extractor(audio, return_tensors="pt", sampling_rate=16000)
+        >>> with torch.no_grad():
+        ...     outputs = model(inputs.input_values, output_hidden_states=True)
         >>> last_hidden = outputs.last_hidden_state  # (1, 50, 768)
         >>> all_hidden = outputs.hidden_states  # Tuple of 13 tensors
         >>> # Or use dict-style access:
@@ -495,6 +500,14 @@ class DistilledSpeechModel(PreTrainedModel):
         )
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         
+        # Ensure input is float (audio samples should be float, not integer tokens)
+        if input_values.dtype in (torch.long, torch.int, torch.int32, torch.int64):
+            input_values = input_values.float()
+        
+        # Ensure 2D input (batch_size, sequence_length)
+        if input_values.dim() == 1:
+            input_values = input_values.unsqueeze(0)
+        
         # Conv encoder: (B, T) -> (B, T', conv_dim)
         extract_features = self.conv_encoder(input_values)
         

preprocessor_config.json

@@ -1,6 +1,7 @@
 {
+  "feature_extractor_type": "Wav2Vec2FeatureExtractor",
   "sampling_rate": 16000,
   "do_normalize": true,
-  "return_attention_mask": false,
-  "feature_extractor_type": "DistilledSpeechFeatureExtractor"
+  "padding_value": 0.0,
+  "return_attention_mask": false
 }