up

app.py

@@ -1,19 +1,41 @@
 from speechbox import PunctuationRestorer
-from datasets import load_dataset
+import soundfile as sf
+import subprocess
+import gradio as gr
 
-streamed_dataset = load_dataset("librispeech_asr", "clean", split="validation", streaming=True)
+restorer = PunctuationRestorer.from_pretrained("openai/whisper-tiny.en")
 
-# get first sample
-sample = next(iter(streamed_dataset))
 
-# print out normalized transcript
-print(sample["text"])
-# => "HE WAS IN A FEVERED STATE OF MIND OWING TO THE BLIGHT HIS WIFE'S ACTION THREATENED TO CAST UPON HIS ENTIRE FUTURE"
+def convert_to_wav(path):
+    if path[-3:] != 'wav':
+        new_path = '.'.join(path.split('.')[:-1]) + '.wav'
+    try:
+        subprocess.call(['ffmpeg', '-i', path, new_path, '-y'])
+    except:  # noqa: E722
+        return path, 'Error: Could not convert file to .wav'
+    path = new_path
+    return path, None
 
-# load the restoring class
-restorer = PunctuationRestorer.from_pretrained("openai/whisper-tiny.en")
-restorer.to("cuda")
 
-restored_text, log_probs = restorer(sample["audio"]["array"], sample["text"], sampling_rate=sample["audio"]["sampling_rate"], num_beams=1)
+def restore(audio, original_transcript):
+    path, error = convert_to_wav(audio)
+    print(error)
+    data, samplerate = sf.read(path)
+
+    text, probs = restorer(data, original_transcript, samplerate, num_beams=1)
+
+    return text, probs
+
 
-print("Restored text:\n", restored_text)
+gr.Interface(
+    title='Punctuation Restorer',
+    fn=restore,
+    inputs=[
+        gr.inputs.Audio(source="upload", type="filepath"),
+        gr.inputs.Text(default="", label="normalized text")
+    ],
+    outputs=[
+        gr.outputs.Textbox(label='Restored text'),
+        gr.outputs.Number(label='Log prob')
+    ]
+  ).launch()

requirements.txt

@@ -1,4 +1,4 @@
 transformers>=4.25.0
 torch
-speechbox
-torchaudio
+speechbox>=0.1.0
+soundfile