added a batch_size input and allow for multiple outputs, also edited the details screen

app.py

@@ -8,7 +8,7 @@ import time
 import gc
 import gradio as gr
 import librosa
-from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM, AutoModelForSeq2SeqLM, AutoTokenizer, AutoProcessor
+from transformers import Wav2Vec2ForCTC, AutoProcessor
 from huggingface_hub import hf_hub_download
 from torchaudio.models.decoder import ctc_decoder
 from numba import cuda
@@ -47,7 +47,7 @@ beam_search_decoder = ctc_decoder(
     tokens=token_file,
     lm=lm_file,
     nbest=1,
-    beam_size=500,
+    beam_size=400,
     beam_size_token=50,
     lm_weight=float(decoding_config["lmweight"]),
     word_score=float(decoding_config["wordscore"]),
@@ -67,7 +67,7 @@ def preprocessAudio(audioFile):
     os.system(f"ffmpeg -y -i {audioFile.name} -ar 16000 ./audioToConvert.wav")
 
 #Transcribe!!!
-def Transcribe(file):
+def Transcribe(file, batch_size):
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     start_time = time.time()
     model.load_adapter("amh")
@@ -75,7 +75,6 @@ def Transcribe(file):
 
     preprocessAudio(file)
     block_size = 30
-    batch_size = 8  # or whatever number you choose
 
     transcripts = []
     speech_segments = []
@@ -94,9 +93,7 @@ def Transcribe(file):
     encoding_start = 0
     encoding_end = 0
     sbv_file = open("subtitle.sbv", "w")
-
-    # Define batch size
-    batch_size = 11
+    transcription_file = open("transcription.txt", "w")
 
     # Create an empty list to hold batches
     batch = []
@@ -122,7 +119,6 @@ def Transcribe(file):
             # Transcribe each segment in the batch
             for i in range(batch_size):
                 transcription = " ".join(beam_search_result[i][0].words).strip()
-                print(transcription)
                 transcripts.append(transcription)
 
                 encoding_end = encoding_start + block_size
@@ -176,15 +172,19 @@ def Transcribe(file):
 
     # Join all transcripts into a single transcript
     transcript = ' '.join(transcripts)
+    transcription_file.write(f"{transcript}")
     sbv_file.close()
+    transcription_file.close()
 
     end_time = time.time()
     print(f"The script ran for {end_time - start_time} seconds.")
-    return("./subtitle.sbv")
+    return(["./subtitle.sbv", "./transcription.txt"])
     
-demo = gr.Interface(fn=Transcribe, inputs=gr.File(label="Upload an audio file of Amharic content"), outputs=gr.File(label="Download .sbv transcription"),
-                   title="Amharic Audio Transcription",
-                    description="This application uses Meta MMS and a custom kenLM model to transcribe Amharic Audio files of arbitrary length into .sbv files. Upload an Amharic audio file and get your transcription! \n(Note: This is only a rough implementation of Meta's MMS for audio transcription, you should manually edit files after transcription has completed.)"
-                   )
+demo = gr.Interface(fn=Transcribe, inputs=[gr.File(label="Upload an audio file of Amharic content"), gr.Slider(0, 25, value=4, step=1, label="batch size", info="Approximately .5GB per batch")], 
+                    outputs=gr.File(label="Download .sbv transcription", file_count="multiple"),
+                    title="Amharic Audio Transcription",
+                    description="This application uses Meta MMS and an Amharic kenLM model to transcribe Amharic Audio files of arbitrary length into .sbv and .txt files. Upload an Amharic audio file and get your transcription! \n(Note: Transcription quality is quite low, you should review and edit transcriptions before making them publicly available)"
+                    )
+
 demo.launch()