NealCaren/transcript

1import whisper2import pandas as pd3import whisper4import subprocess5from simple_diarizer.diarizer import Diarizer6import streamlit as st7import base648import tempfile9 10 11 12def create_download_link(val, filename, label):13    '''Hack to have a stable download link in Streamlit'''14    b64 = base64.b64encode(val)15    return f'<a href="data:application/octet-stream;base64,{b64.decode()}" download="{filename}">{label}</a>'16 17 18def segment(nu_speakers):19    '''Segment the audio using simple_diarizer.20    Defaults to the speechbrain ECAPA-TDNN embeddings.'''21 22    diar = Diarizer(embed_model='ecapa',cluster_method='sc')23    segments = diar.diarize(temp_file, num_speakers=nu_speakers)24 25    sdf = pd.DataFrame(segments)26 27    # reorganize so the first speaker is always speaker 128    speaker_s = sdf['label'].drop_duplicates().reset_index()['label']29    speaker_d = dict((v,k+1) for k,v in speaker_s.items())30 31    sdf['speaker'] = sdf['label'].replace(speaker_d)32    return sdf33def monotize(uploaded):34    '''Convert the upload file to audio file.'''35    cmd = f"ffmpeg -y -i {uploaded} -acodec pcm_s16le -ar 16000 -ac 1 {temp_file}"36    subprocess.Popen(cmd, shell=True).wait()37 38def audio_to_df(uploaded):39    '''Turn the upload file in a segemented dataframe.'''40    #monotize(uploaded)41    model = whisper.load_model(model_size)42    result = model.transcribe(temp_file,43                              without_timestamps=False,44                              task = task)45    tdf = pd.DataFrame(result['segments'])46    return tdf47 48 49 50def add_preface(row):51    ''' Add speaker prefix to transcript during transcribe().'''52    text = row['text'].replace('\n','')53    speaker = row['speaker']54    return f'Speaker {speaker}: {text}'55 56def transcribe(uploaded, nu_speakers):57    # Convert file to mono58    with st.spinner(text="Converting file..."):59        monotize('temp_audio')60 61    # Make audio available to play in UI62    audio_file = open(temp_file, 'rb')63    audio_bytes = audio_file.read()64    st.audio(temp_file, format='audio/wav')65 66    # trancibe file67    with st.spinner(text=f"Transcribing using {model_size} model..."):68        tdf = audio_to_df(uploaded)69    # segement file70    with st.spinner(text="Segmenting..."):71        sdf = segment(nu_speakers)72 73    # Find the nearest transcript line to the start of each speaker74    ns_list = sdf[['start','speaker']].to_dict(orient='records')75    for row in ns_list:76        input = row['start']77        id = tdf.iloc[(tdf['start']-input).abs().argsort()[:1]]['id'].values[0]78        tdf.loc[tdf['id'] ==id, 'speaker'] = row['speaker']79    tdf['speaker'].fillna(method = 'ffill', inplace = True)80    tdf['speaker'].fillna(method = 'bfill', inplace = True)81    tdf['n1'] = tdf['speaker'] != tdf['speaker'].shift(1)82    tdf['speach'] = tdf['n1'].cumsum()83 84    # collaps the dataframe by speach turn.85    binned_df = tdf.groupby(['speach', 'speaker'])['text'].apply('\n'.join).reset_index()86    binned_df['speaker'] = binned_df['speaker'].astype(int)87    binned_df['output'] = binned_df.apply(add_preface, axis=1)88 89    # Display the transcript and prepare for export90    lines = []91    for row in binned_df['output'].values:92        st.write(row)93        lines.append(row)94    tdf['speaker'] = tdf['speaker'].astype(int)95 96    tdf_cols = ['speaker','start','end','text']97    #st.dataframe(tdf[tdf_cols])98    return {'text':lines, 'df': tdf[tdf_cols]}99 100 101descript = ("This web app creates transcripts using OpenAI's [Whisper](https://github.com/openai/whisper) to transcribe "102            "audio files combined with [Chau](https://github.com/cvqluu)'s [Simple Diarizer](https://github.com/cvqluu/simple_diarizer) "103            "to partition the text by speaker.\n"104            "* You can upload an audio or video file of up to 200MBs.\n"105            "* Creating the transcript takes some time. "106            "The process takes approximately 20% of the length of the audio file using the base Whisper model.\n "107            "* The transcription process handles a variety of languages, and can also translate the audio to English. The tiny model is not good at translating. \n"108            "* Speaker segmentation seems to work best with the base model. The small model produces better transcripts, but something seems off with the timecodes, degrading the speaker attribution. \n"109            "* After uploading the file, be sure to select the number of speakers." )110 111st.title("Automated Transcription")112st.markdown(descript)113 114form = st.form(key='my_form')115uploaded = form.file_uploader("Choose a file")116nu_speakers = form.slider('Number of speakers in recording:', min_value=1, max_value=8, value=2, step=1)117models = form.selectbox(118    'Which Whisper model?',119    ('Tiny (fast)', 'Base (good)', 'Small (great but slow)', 'Medium (greater but slower)'), index=1)120translate = form.checkbox('Translate to English?')121submit = form.form_submit_button("Transcribe!")122 123 124if submit:125    if models == 'Tiny (fast)':126        model_size = 'tiny'127    elif models == 'Base (good)':128        model_size ='base'129    elif models == 'Small (great but slow)':130        model_size = 'small'131    elif models == 'Medium (greater but slower)':132        model_size = 'medium'133 134    if translate == True:135        task = 'translate'136    else:137        task = 'transcribe'138 139    #temporary file to store audio_file140    tmp_dir = tempfile.TemporaryDirectory()141    temp_file = tmp_dir.name + '/mono.wav'142 143    bytes_data = uploaded.getvalue()144    with open('temp_audio', 'wb') as outfile:145        outfile.write(bytes_data)146 147 148    # Transcribe/translate and segment149    transcript = transcribe('temp_audio', nu_speakers)150 151    # Prepare text file for export.152    text = '\n'.join(transcript['text']).encode('utf-8')153    download_url = create_download_link(text, 'transcript.txt', 'Download transcript as plain text.')154    st.markdown(download_url, unsafe_allow_html=True)155 156    # prepare CSV file for expport.157    csv = transcript['df'].to_csv( float_format='%.2f', index=False).encode('utf-8')158    download_url = create_download_link(csv, 'transcript.csv', 'Download transcript as CSV (with time codes)')159    st.markdown(download_url, unsafe_allow_html=True)160    tmp_dir.cleanup()161