note-taker / app.py
msaelices's picture
Allow to customize the whisper model
ab9ec7c
raw
history blame contribute delete
No virus
3.73 kB
import os
import streamlit as st
from dotenv import load_dotenv
from engines import get_engine
import api
# Load environment variables from .env file before importing any other modules
load_dotenv()
def main():
st.set_page_config(
page_title='Note Taker',
page_icon='πŸŽ™οΈ',
layout='centered',
initial_sidebar_state='expanded',
)
title = 'πŸŽ™οΈ Meetings Note Taker πŸŽ™οΈ'
st.title(title)
st.write(
'Upload an audio file, transcribe it using WhisperX or AssemblyAI, and generate meeting notes using your selected model.'
)
openai_api_key = os.environ.get('OPENAI_API_KEY') or st.text_input(
'Enter your OpenAI API key:', type='password'
)
engine_type = os.environ.get('TRANSCRIPTION_ENGINE') or st.selectbox(
'Select a transcription engine:', ['WhisperX', 'AssemblyAI']
)
if engine_type == 'WhisperX':
device = os.environ.get('PYTORCH_DEVICE') or st.selectbox(
'Select a PyTorch device:', ['cuda', 'cpu']
)
compute_type = os.environ.get('PYTORCH_COMPUTE_TYPE') or st.selectbox(
'Select a compute type:', ['int8', 'float16']
)
batch_size = os.environ.get('PYTORCH_BATCH_SIZE') or st.selectbox(
'Select a batch size:', [4, 8, 16, 32, 64]
)
whisper_model = os.environ.get('WHISPER_MODEL') or st.selectbox(
'Select a Whisper model:', ['large-v2', 'base']
)
else:
device = None
compute_type = None
batch_size = None
whisper_model = None
engine_api_key = os.environ.get(
f'{engine_type.upper()}_API_KEY'
) or st.text_input(f'Enter your {engine_type} API key:', type='password')
openai_model = os.environ.get('OPENAI_MODEL') or st.selectbox(
'Select a model:', ['gpt-3.5-turbo-16k', 'gpt-4-0613']
)
uploaded_audio = st.file_uploader(
'Upload an audio file',
type=['aac', 'm4a', 'mp3', 'webm', 'mp4', 'mpga', 'wav', 'mpeg'],
accept_multiple_files=False,
)
language = os.environ.get('AUDIO_LANGUAGE') or st.selectbox(
'Language code of the audio:', ['en', 'es']
)
prompt = os.environ.get('SUMMARIZE_PROMPT') or st.text_area(
'Enter a custom prompt for the summary:', 'Summarize the following audio transcription with a list of the key points with the speakers in the original language'
)
if st.button('Generate Notes'):
if uploaded_audio:
if openai_api_key:
st.markdown('Transcribing the audio...')
engine = get_engine(
engine_type,
api_key=engine_api_key,
device=device,
compute_type=compute_type,
batch_size=batch_size,
whisper_model=whisper_model,
)
transcription = api.transcribe(engine, language, uploaded_audio)
st.markdown(
f'### Transcription:\n\n<details><summary>Click to view</summary><p><pre><code>{transcription}</code></pre></p></details>',
unsafe_allow_html=True,
)
st.markdown('Summarizing the transcription...')
summary = api.summarize_transcript(
openai_api_key=openai_api_key,
transcript=transcription,
openai_model=openai_model,
prompt=prompt,
)
st.markdown(f'### Summary:')
st.write(summary)
else:
st.error('We need valid OpenAI API keys')
if __name__ == '__main__':
main()