Update audio_tokenizer.py

audio_tokenizer.py

@@ -124,31 +124,54 @@ def process_dataset(
         text_field: Name of text field in dataset (default: "text_scribe")
         target_sample_rate: Target audio sample rate (default: 24000)
     """
-    # Set tokenizer and config based on model type
-    if model_type == "qwen3":
-        tokenizer_model = "Qwen/Qwen3-0.6B"
-        config_path = "qwen3.yaml"
-    elif model_type == "lfm2":
-        tokenizer_model = "LiquidAI/LFM2-350M"
-        config_path = "lfm2.yaml"
-    else:
-        raise ValueError(f"Invalid model_type: {model_type}. Must be 'qwen3' or 'lfm2'")
-
-    # Load configuration
-    print(f"Loading config from: {config_path}")
-    config = load_config(config_path)
-
-    TOKENIZER_LENGTH = config['TOKENIZER_LENGTH']
-    START_OF_TEXT = config['START_OF_TEXT']
-    END_OF_TEXT = config['END_OF_TEXT']
-    START_OF_SPEECH = config['START_OF_SPEECH']
-    END_OF_SPEECH = config['END_OF_SPEECH']
-    START_OF_HUMAN = config['START_OF_HUMAN']
-    END_OF_HUMAN = config['END_OF_HUMAN']
-    START_OF_AI = config['START_OF_AI']
-    END_OF_AI = config['END_OF_AI']
-    PAD_TOKEN = config['PAD_TOKEN']
-    AUDIO_TOKENS_START = config['AUDIO_TOKENS_START']
+    try:
+        # Set tokenizer and config based on model type
+        print(f"Setting up configuration for model_type: {model_type}")
+        if model_type == "qwen3":
+            tokenizer_model = "Qwen/Qwen3-0.6B"
+            config_path = "qwen3.yaml"
+        elif model_type == "lfm2":
+            tokenizer_model = "LiquidAI/LFM2-350M"
+            config_path = "lfm2.yaml"
+        else:
+            raise ValueError(f"Invalid model_type: {model_type}. Must be 'qwen3' or 'lfm2'")
+
+        print(f"Tokenizer Model: {tokenizer_model}")
+        print(f"Config Path: {config_path}")
+
+        # Load configuration
+        print(f"Loading config from: {config_path}")
+        config = load_config(config_path)
+        print(f"Config loaded successfully. Type: {type(config)}")
+
+        if not isinstance(config, dict):
+            raise TypeError(f"Config must be a dictionary, got {type(config)}")
+
+    except Exception as e:
+        print(f"Error in initial setup: {str(e)}")
+        raise
+
+    try:
+        print("Extracting config values...")
+        TOKENIZER_LENGTH = config['TOKENIZER_LENGTH']
+        START_OF_TEXT = config['START_OF_TEXT']
+        END_OF_TEXT = config['END_OF_TEXT']
+        START_OF_SPEECH = config['START_OF_SPEECH']
+        END_OF_SPEECH = config['END_OF_SPEECH']
+        START_OF_HUMAN = config['START_OF_HUMAN']
+        END_OF_HUMAN = config['END_OF_HUMAN']
+        START_OF_AI = config['START_OF_AI']
+        END_OF_AI = config['END_OF_AI']
+        PAD_TOKEN = config['PAD_TOKEN']
+        AUDIO_TOKENS_START = config['AUDIO_TOKENS_START']
+        print("✓ All config values extracted successfully")
+    except KeyError as e:
+        print(f"Missing key in config: {e}")
+        print(f"Available keys: {list(config.keys())}")
+        raise
+    except Exception as e:
+        print(f"Error extracting config values: {str(e)}")
+        raise
 
     # Download dataset
     print(f"Downloading dataset: {original_dataset}")