Update index.js

index.js

@@ -1,4 +1,4 @@
-import { pipeline } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.6';
+import { WhisperForConditionalGeneration, WhisperProcessor } from 'https://cdn.jsdelivr.net/npm/@huggingface/transformers@3.7.6';
 
 // Get DOM elements
 const status = document.getElementById('status');
@@ -13,7 +13,8 @@ const recordingTimeDisplay = document.getElementById('recordingTime');
 const visualizerBars = document.querySelectorAll('.bar');
 
 // State
-let transcriber = null;
+let model = null;
+let processor = null;
 let mediaStream = null;
 let audioContext = null;
 let mediaRecorder = null;
@@ -32,45 +33,48 @@ async function initModel() {
         status.className = 'loading';
         
         const device = useWebGPUCheckbox.checked ? 'webgpu' : 'wasm';
+        const dtype = useWebGPUCheckbox.checked ? 'fp32' : 'fp32';
         
-        // Load your custom ATOM model
-        // IMPORTANT: This automatically loads the custom tokenizer from Chillarmo/ATOM
-        // The tokenizer is specific to your model's training and must be used
-        transcriber = await pipeline(
-            'automatic-speech-recognition',
-            'Chillarmo/ATOM',
-            {
-                device: device,
-                // Ensure we use the model's own tokenizer, not a default one
-                revision: 'main', // Use main branch which has your custom tokenizer
-                progress_callback: (progress) => {
-                    if (progress.status === 'downloading') {
-                        const percent = Math.round((progress.loaded / progress.total) * 100);
-                        status.textContent = `Downloading ${progress.file}: ${percent}%`;
-                    } else if (progress.status === 'loading') {
-                        status.textContent = `Loading ${progress.file}...`;
-                    } else if (progress.status === 'progress') {
-                        const percent = Math.round(progress.progress);
-                        status.textContent = `Loading model: ${percent}%`;
-                    }
+        // Load processor (includes the custom Armenian tokenizer)
+        status.textContent = 'Loading custom Armenian processor/tokenizer...';
+        processor = await WhisperProcessor.from_pretrained('Chillarmo/ATOM', {
+            progress_callback: (progress) => {
+                if (progress.status === 'downloading') {
+                    const percent = Math.round((progress.loaded / progress.total) * 100);
+                    status.textContent = `Downloading ${progress.file}: ${percent}%`;
                 }
             }
-        );
+        });
+        
+        console.log('✓ ATOM Processor loaded (includes custom tokenizer)');
         
-        console.log('ATOM model loaded successfully with custom Armenian tokenizer');
-        status.textContent = 'Model & custom tokenizer loaded! Ready to transcribe Armenian speech.';
+        // Load model
+        status.textContent = 'Loading ATOM model...';
+        model = await WhisperForConditionalGeneration.from_pretrained('Chillarmo/ATOM', {
+            device: device,
+            dtype: dtype,
+            progress_callback: (progress) => {
+                if (progress.status === 'downloading') {
+                    const percent = Math.round((progress.loaded / progress.total) * 100);
+                    status.textContent = `Downloading model ${progress.file}: ${percent}%`;
+                } else if (progress.status === 'loading') {
+                    status.textContent = `Loading ${progress.file}...`;
+                }
+            }
+        });
+        
+        console.log('✓ ATOM Model loaded');
+        console.log('Model config:', model.config);
+        console.log('Processor:', processor);
+        
+        status.textContent = 'ATOM ready! Model + custom tokenizer loaded successfully.';
         status.className = 'ready';
         startBtn.disabled = false;
     } catch (error) {
         console.error('Model loading error:', error);
         status.textContent = `Error loading model: ${error.message}`;
         status.className = 'error';
-        
-        // Log more details for debugging tokenizer issues
         console.error('Full error details:', error);
-        if (error.message.includes('tokenizer')) {
-            status.textContent = 'Error: Custom tokenizer failed to load. Check console.';
-        }
     }
 }
 
@@ -232,16 +236,33 @@ async function processAudioChunk(chunks) {
         // Get audio data as Float32Array
         const audioData = audioBuffer.getChannelData(0);
         
-        // Transcribe with ATOM model
-        // This uses the custom Armenian tokenizer that comes with the model
-        // The tokenizer is critical - it was trained with the model
-        const result = await transcriber(audioData, {
+        console.log('Processing audio chunk:', audioData.length, 'samples at', audioBuffer.sampleRate, 'Hz');
+        
+        // Process audio with the processor (includes custom tokenizer)
+        const inputs = await processor(audioData, {
             sampling_rate: audioBuffer.sampleRate,
         });
         
+        console.log('Processor output:', inputs);
+        
+        // Generate with the model
+        const outputs = await model.generate({
+            ...inputs,
+        });
+        
+        console.log('Model outputs:', outputs);
+        
+        // Decode the output tokens using the custom tokenizer
+        const decoded = processor.batch_decode(outputs, {
+            skip_special_tokens: true,
+        });
+        
+        console.log('Decoded text:', decoded);
+        
         // Add to transcription
-        if (result && result.text && result.text.trim()) {
-            addTranscription(result.text.trim());
+        const text = decoded[0].trim();
+        if (text) {
+            addTranscription(text);
             chunkCount++;
             chunkCountDisplay.textContent = chunkCount;
         }
@@ -260,6 +281,7 @@ async function processAudioChunk(chunks) {
         console.error('Error processing audio:', error);
         status.textContent = `Processing error: ${error.message}`;
         status.className = 'error';
+        console.error('Full processing error:', error);
         
         // Restore recording status if still recording
         setTimeout(() => {