Update app.py

app.py

@@ -1,73 +1,80 @@
 import os
+# disable TorchDynamo since UnsloTh models can have issues with TorchDynamo
 os.environ["TORCHDYNAMO_DISABLE"] = "1"
 
 import gradio as gr
 import torch
 import librosa
 import numpy as np
-import glob
+from threading import Thread
 from unsloth import FastModel
 from transformers import AutoProcessor, TextIteratorStreamer
-from threading import Thread
 
 TARGET_SAMPLING_RATE = 16000
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-print("Loading processor and model...")
+print("Loading model + processor...")
 
+# load the processor & model from the right repo
 processor = AutoProcessor.from_pretrained("EpistemeAI/Audiogemma-3N-finetune")
-
 model, _ = FastModel.from_pretrained(
     model_name="EpistemeAI/Audiogemma-3N-finetune",
     dtype=None,
-    max_seq_length=1024,
+    max_seq_length=2048,
     load_in_4bit=True,
     full_finetuning=False,
+    device_map="auto"
 )
 
 model.eval()
-
-print("Model loaded on", device)
+print("Loaded Gemma-3N on", device)
 
 def transcribe_and_translate(audio_input):
     if audio_input is None:
-        yield "Please upload or record an audio file."
+        yield "Upload or record audio first."
         return
 
     sample_rate, audio_array = audio_input
 
+    # mono
     if audio_array.ndim > 1:
         audio_array = audio_array.mean(axis=1)
 
     audio_array = audio_array.astype(np.float32)
 
+    # resample to 16k
     if sample_rate != TARGET_SAMPLING_RATE:
         audio_array = librosa.resample(
             audio_array, orig_sr=sample_rate, target_sr=TARGET_SAMPLING_RATE
         )
 
+    # prepare prompt
     messages = [
         {
             "role": "system",
             "content": [
-                {"type": "text", "text": "You are an assistant that transcribes and translates speech accurately."}
+                {
+                    "type": "text",
+                    "text": "You are a model that accurately transcribes spoken audio and translates it to German."
+                }
             ],
         },
         {
             "role": "user",
             "content": [
                 {"type": "audio", "audio": audio_array},
-                {"type": "text", "text": "Please transcribe this audio and translate it to German."}
+                {"type": "text", "text": "Transcribe the spoken audio and translate to German."}
             ],
         },
     ]
 
+    # tokenize & prep inputs
     inputs = processor.apply_chat_template(
         messages,
         add_generation_prompt=True,
         tokenize=True,
         return_dict=True,
-        return_tensors="pt",
+        return_tensors="pt"
     ).to(device)
 
     streamer = TextIteratorStreamer(processor, skip_prompt=True)
@@ -75,46 +82,39 @@ def transcribe_and_translate(audio_input):
     generation_kwargs = dict(
         **inputs,
         max_new_tokens=1024,
-        temperature=0.8,
+        temperature=1.0,
         top_p=0.95,
-        top_k=64,
-        streamer=streamer,
+        top_k=50,
+        streamer=streamer
     )
 
     thread = Thread(target=model.generate, kwargs=generation_kwargs)
     thread.start()
 
     output_text = ""
-    for new_text in streamer:
-        output_text += new_text
+    for chunk in streamer:
+        output_text += chunk
         yield output_text
 
-
-example_audios = glob.glob("test_wav_files/*.wav")
-example_list = [[audio] for audio in example_audios]
-
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown(
-        """
-        # Audio Transcription & Translation (Gemma-3N)
-        Upload or record audio and receive transcription and German translation.
-        Powered by Audiogemma-3N + Unsloth.
-        """
-    )
+    gr.Markdown("# Gemma-3N Audio Transcription + German Translation")
 
     with gr.Row():
-        audio_input = gr.Audio(sources=["upload", "microphone"], type="numpy", label="Audio Input")
-        text_output = gr.Textbox(label="Transcription & Translation", lines=12)
-
-    submit_btn = gr.Button("Transcribe and Translate", variant="primary")
-    submit_btn.click(fn=transcribe_and_translate, inputs=audio_input, outputs=text_output)
+        audio_input = gr.Audio(
+            sources=["upload","microphone"],
+            type="numpy",
+            label="Your Audio"
+        )
+        text_output = gr.Textbox(
+            label="Transcript & Translation",
+            lines=10
+        )
 
-    gr.Examples(
-        examples=example_list,
-        inputs=audio_input,
-        outputs=text_output,
+    submit_btn = gr.Button("Transcribe + Translate")
+    submit_btn.click(
         fn=transcribe_and_translate,
-        cache_examples=False,
+        inputs=audio_input,
+        outputs=text_output
     )
 
 if __name__ == "__main__":