File size: 4,622 Bytes
a95e64c
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4e9010c
a95e64c
 
 
 
 
 
da7f7e0
a95e64c
 
da7f7e0
 
a95e64c
 
da7f7e0
 
a95e64c
4e9010c
a95e64c
 
 
 
 
 
da7f7e0
a95e64c
 
 
09330c2
a95e64c
 
 
 
 
da7f7e0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
09330c2
da7f7e0
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
4e9010c
a95e64c
 
da7f7e0
 
a95e64c
da7f7e0
 
 
 
 
 
 
a95e64c
 
 
 
da7f7e0
a95e64c
 
 
da7f7e0
a95e64c
 
 
 
 
 
1db1226
a95e64c
4890ad0
 
 
 
a95e64c
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import soundfile as sf
import datetime
from pyctcdecode import BeamSearchDecoderCTC
import torch
import os
import time
import gc
import gradio as gr
import librosa
from transformers import Wav2Vec2ForCTC, Wav2Vec2ProcessorWithLM, AutoModelForSeq2SeqLM, AutoTokenizer
from numba import cuda

# load pretrained model
model = Wav2Vec2ForCTC.from_pretrained("facebook/mms-1b-all")
processor = Wav2Vec2ProcessorWithLM.from_pretrained("jlonsako/mms-1b-all-AmhLM")


#Define Functions

#convert time into .sbv format
def format_time(seconds):
    # Convert seconds to hh:mm:ss,ms format
    return str(datetime.timedelta(seconds=seconds)).replace('.', ',')

#Convert Video/Audio into 16K wav file
def preprocessAudio(audioFile):
    os.system(f"ffmpeg -y -i {audioFile.name} -ar 16000 ./audioToConvert.wav")

#Transcribe!!!
def Transcribe(file):
    device = "cuda:0" if torch.cuda.is_available() else "cpu"
    start_time = time.time()
    model.load_adapter("amh")
    model.half()

    preprocessAudio(file)
    block_size = 30
    batch_size = 22  # or whatever number you choose

    transcripts = []
    speech_segments = []

    stream = librosa.stream(
        "./audioToConvert.wav",
        block_length=block_size,
        frame_length=16000,
        hop_length=16000
    )

    model.to(device)
    print("Model loaded to gpu: Entering transcription phase")

    #Code for timestamping
    encoding_start = 0
    encoding_end = 0
    sbv_file = open("subtitle.sbv", "w")

    for speech_segment in stream:  
        if len(speech_segment.shape) > 1:
            speech_segment = speech_segment[:,0] + speech_segment[:,1]
        speech_segments.append(speech_segment)

        if len(speech_segments) == batch_size:
            input_values = processor(speech_segments, sampling_rate=16_000, return_tensors="pt", padding=True).input_values.to(device)
            input_values = input_values.half()
            with torch.no_grad():
                logits = model(input_values).logits
            if len(logits.shape) == 1:
                logits = logits.unsqueeze(0)
            #predicted_ids = torch.argmax(logits, dim=-1)
            transcriptions = processor.batch_decode(logits.cpu().numpy()).text
            transcripts.extend(transcriptions)

            # Write to the .sbv file
            for i, transcription in enumerate(transcriptions):
                encoding_start = encoding_end  # Maintain the 'encoding_start' across batches
                encoding_end = encoding_start + block_size
                formatted_start = format_time(encoding_start)
                formatted_end = format_time(encoding_end)
                sbv_file.write(f"{formatted_start},{formatted_end}\n")
                sbv_file.write(f"{transcription}\n\n")

            # Clear the batch
            speech_segments = []

            # Freeing up memory
            del input_values
            del logits
            del transcriptions
            torch.cuda.empty_cache()
            gc.collect()

    if speech_segments:
        input_values = processor(speech_segments, sampling_rate=16_000, return_tensors="pt", padding=True).input_values.to(device)
        input_values = input_values.half()
        with torch.no_grad():
            logits = model(input_values).logits
        transcriptions = processor.batch_decode(logits.cpu().numpy()).text
        transcripts.extend(transcriptions)

        for i in range(len(speech_segments)):
            encoding_end = encoding_start + block_size
            formatted_start = format_time(encoding_start)
            formatted_end = format_time(encoding_end)
            sbv_file.write(f"{formatted_start},{formatted_end}\n")
            sbv_file.write(f"{transcriptions[i]}\n\n")
            encoding_start = encoding_end

        # Freeing up memory
        del input_values
        del logits
        del transcriptions
        torch.cuda.empty_cache()
        gc.collect()


    # Join all transcripts into a single transcript
    transcript = ' '.join(transcripts)
    sbv_file.close()

    end_time = time.time()
    print(f"The script ran for {end_time - start_time} seconds.")
    return("./subtitle.sbv")
    
demo = gr.Interface(fn=Transcribe, inputs=gr.File(label="Upload an audio file of Amharic content"), 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!"
                   )
demo.launch()