Update app.py

app.py

@@ -10,53 +10,75 @@ model = WhisperForConditionalGeneration.from_pretrained(model_id)
 device = "cuda" if torch.cuda.is_available() else "cpu"
 model.to(device)
 
-forced_decoder_ids = processor.get_decoder_prompt_ids(language="he", task="transcribe")
+# Force Hebrew (transcribe) decoding:
+forced_decoder_ids = processor.get_decoder_prompt_ids(
+    language="he", 
+    task="transcribe"
+)
+
+stop_processing = False
+def stop():
+    global stop_processing
+    stop_processing = True
 
-def transcribe_long(audio_file):
-    # 1) Load full audio (limit to 60 minutes)
+def transcribe_first_chunk(audio_file):
+    """
+    Transcribe only the first 'time_limit_s' seconds of the uploaded audio.
+    """
+    global stop_processing
+    stop_processing = False
+
+    # A) Load at 16kHz
     waveform, sr = librosa.load(audio_file, sr=16000)
-    if len(waveform) > sr * 3600:
-        waveform = waveform[: sr * 3600]
-
-    # 2) Split into ~2min chunks
-    chunk_sec = 120
-    chunk_size = sr * chunk_sec
-    all_text = []
-    for start in range(0, len(waveform), chunk_size):
-        chunk = waveform[start : start + chunk_size]
-        # skip chunks <2s if you want
-        if len(chunk) < sr * 2:
-            continue
-
-        # 3) Encode with attention mask
-        inputs = processor(
-            chunk,
-            sampling_rate=16000,
-            return_tensors="pt",
-            padding="longest",
-            return_attention_mask=True
+
+    # B) Truncate to the first 4 minutes
+    time_limit_s = 4 * 60  # 4 minutes = 240 seconds
+    if len(waveform) > sr * time_limit_s:
+        waveform = waveform[: sr * time_limit_s]
+
+    # Also limit if total is over 60 min (safety)
+    max_audio_sec = 60 * 60
+    if len(waveform) > sr * max_audio_sec:
+        waveform = waveform[: sr * max_audio_sec]
+
+    # C) Preprocess: get attention mask
+    inputs = processor(
+        waveform,
+        sampling_rate=16000,
+        return_tensors="pt",
+        padding="longest",
+        return_attention_mask=True
+    )
+    input_features = inputs.input_features.to(device)
+    attention_mask = inputs.attention_mask.to(device)
+
+    if stop_processing:
+        return "⚠️ Stopped by User ⚠️"
+
+    # D) Generate
+    with torch.no_grad():
+        predicted_ids = model.generate(
+            input_features,
+            attention_mask=attention_mask,
+            max_new_tokens=444,          # keep total under 448 tokens
+            do_sample=False,             # deterministic
+            forced_decoder_ids=forced_decoder_ids  # ensure Hebrew transcription
         )
-        input_features = inputs.input_features.to(device)
-        attention_mask = inputs.attention_mask.to(device)
-
-        # 4) Generate
-        with torch.no_grad():
-            predicted_ids = model.generate(
-                input_features,
-                attention_mask=attention_mask,
-                max_new_tokens=444,
-                do_sample=False,
-                forced_decoder_ids=forced_decoder_ids
-            )
-        text_chunk = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
-        all_text.append(text_chunk)
-
-    return " ".join(all_text)
-
-demo = gr.Interface(
-    fn=transcribe_long,
-    inputs=gr.Audio(type="filepath", label="Upload Audio (unlimited)"),
-    outputs="text",
-    title="Chunked Whisper (No Token Overflow)"
-)
+
+    # E) Decode
+    text = processor.batch_decode(predicted_ids, skip_special_tokens=True)[0]
+    return text
+
+with gr.Blocks() as demo:
+    gr.Markdown("## Hebrew Whisper (Only First 4 Minutes)")
+
+    audio_input = gr.Audio(type="filepath", label="Upload Audio (Truncate to 4min)")
+    output_text = gr.Textbox(label="Partial Transcription")
+
+    start_btn = gr.Button("Start Transcription")
+    stop_btn = gr.Button("Stop Processing", variant="stop")
+
+    start_btn.click(transcribe_first_chunk, inputs=audio_input, outputs=output_text)
+    stop_btn.click(stop)
+
 demo.launch()