Update app.py

app.py

@@ -7,63 +7,35 @@ import librosa
 import math
 from transformers import MoonshineForConditionalGeneration, AutoProcessor
 
-# Use GPU if available and set appropriate dtype
 device = "cuda:0" if torch.cuda.is_available() else "cpu"
 torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
-# Load model and processor - Moonshine Tiny
 model = MoonshineForConditionalGeneration.from_pretrained('UsefulSensors/moonshine-tiny').to(device).to(torch_dtype)
 processor = AutoProcessor.from_pretrained('UsefulSensors/moonshine-tiny')
 
-# --- Longer token limits (simple) ---
-TOKENS_PER_SEC = 12.0       # was ~7.0 before
-MIN_NEW_TOKENS = 48         # was 24; gives short clips more room
-MAX_NEW_TOKENS_CAP = 3200   # generous cap to avoid runaway
+TOKENS_PER_SEC = 12.0
+MIN_NEW_TOKENS = 48
+MAX_NEW_TOKENS_CAP = 1600
 
-# Define transcription function using HF Zero GPU
 @spaces.GPU
 def transcribe_audio(audio_file):
     if not audio_file:
         return "No audio provided."
-
-    # Load and preprocess audio
     audio_array, sr = sf.read(audio_file)
     if audio_array.ndim > 1:
-        audio_array = np.mean(audio_array, axis=1) 
-
-    # Resample if necessary in case the audio file has a different sampling rate
+        audio_array = np.mean(audio_array, axis=1)
     target_sr = processor.feature_extractor.sampling_rate
     if sr != target_sr:
         audio_array = librosa.resample(audio_array, orig_sr=sr, target_sr=target_sr)
-
-    # Prepare inputs for the model - ensure correct dtype and device
-    inputs = processor(
-        audio_array,
-        sampling_rate=target_sr,
-        return_tensors="pt"
-    ).to(device, torch_dtype)
-
-    # Duration-based max_new_tokens calculation (longer limits)
+    inputs = processor(audio_array, sampling_rate=target_sr, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}
     duration_sec = len(audio_array) / float(target_sr)
-    max_new_tokens = min(
-        MAX_NEW_TOKENS_CAP,
-        max(MIN_NEW_TOKENS, int(math.ceil(duration_sec * TOKENS_PER_SEC)))
-    )
-
-    # Generate transcription with adjusted max_new_tokens
-    generated_ids = model.generate(**inputs, do_sample=False, max_new_tokens=max_new_tokens)
-    return processor.decode(generated_ids[0], skip_special_tokens=True) # Decode the generated IDs to text
+    max_new_tokens = min(MAX_NEW_TOKENS_CAP, max(MIN_NEW_TOKENS, int(math.ceil(duration_sec * TOKENS_PER_SEC))))
+    generated_ids = model.generate(**inputs, do_sample=False, max_new_tokens=max_new_tokens, no_repeat_ngram_size=4, repetition_penalty=1.05)
+    return processor.decode(generated_ids[0], skip_special_tokens=True)
 
-# Set Gradio theme
-theme = gr.themes.Ocean(
-    primary_hue="indigo",
-    secondary_hue="fuchsia",
-    neutral_hue="slate",
-).set(
-    button_large_radius='*radius_sm'
-)
+theme = gr.themes.Ocean(primary_hue="indigo", secondary_hue="fuchsia", neutral_hue="slate").set(button_large_radius='*radius_sm')
 
-# Create Gradio interface
 with gr.Blocks(theme=theme) as demo:
     gr.Markdown("## Moonshine Tiny STT - 27M Parameters")
     gr.HTML("""
@@ -73,46 +45,17 @@ with gr.Blocks(theme=theme) as demo:
                 alt="VibeVoice Banner">
         </div>
         """)
-    
     with gr.Tabs():
         with gr.TabItem("Upload Audio"):
-            audio_file = gr.Audio(
-                sources=["upload"],
-                type="filepath",
-                label="Upload Audio File"
-            )
-            output_text1 = gr.Textbox(
-                label="Transcription",
-                placeholder="Transcription will appear here...",
-                lines=20,
-                autoscroll=True
-            )
+            audio_file = gr.Audio(sources=["upload"], type="filepath", label="Upload Audio File")
+            output_text1 = gr.Textbox(label="Transcription", placeholder="Transcription will appear here...", lines=20, autoscroll=True)
             upload_button = gr.Button("Transcribe Uploaded Audio")
-            upload_button.click(
-                fn=transcribe_audio,
-                inputs=audio_file,
-                outputs=output_text1
-            )
-            
+            upload_button.click(fn=transcribe_audio, inputs=audio_file, outputs=output_text1)
         with gr.TabItem("Record Audio"):
-            audio_mic = gr.Audio(
-                sources=["microphone"],
-                type="filepath",
-                label="Record Audio"
-            )
-            output_text2 = gr.Textbox(
-                label="Transcription",
-                placeholder="Transcription will appear here...",
-                lines=20,
-                autoscroll=True
-            )
+            audio_mic = gr.Audio(sources=["microphone"], type="filepath", label="Record Audio")
+            output_text2 = gr.Textbox(label="Transcription", placeholder="Transcription will appear here...", lines=20, autoscroll=True)
             record_button = gr.Button("Transcribe Recorded Audio")
-            record_button.click(
-                fn=transcribe_audio,
-                inputs=audio_mic,
-                outputs=output_text2
-            )
-
+            record_button.click(fn=transcribe_audio, inputs=audio_mic, outputs=output_text2)
     gr.Markdown("""
     ### Instructions:
     1. Choose either 'Upload Audio' or 'Record Audio' tab