Upload 2 files

app.py

@@ -0,0 +1,156 @@
+import streamlit as st
+from streamlit_mic_recorder import mic_recorder
+from transformers import pipeline
+import torch
+from transformers import BertTokenizer, BertForSequenceClassification
+import numpy as np
+import pandas as pd
+import time
+import altair as alt
+
+
+def callback():
+    if st.session_state.my_recorder_output:
+        audio_bytes = st.session_state.my_recorder_output['bytes']
+        st.audio(audio_bytes)
+
+
+@st.cache_resource
+def load_text_to_speech_model(model="openai/whisper-base"):
+    pipe = pipeline("automatic-speech-recognition", model=model)
+    return pipe
+
+
+def translate(inputs, model="openai/whisper-base"):
+    pipe = pipeline("automatic-speech-recognition", model=model)
+    translate_result = pipe(inputs, generate_kwargs={'task': 'translate'})
+    return translate_result['text']
+
+
+@st.cache_resource
+def load_classification_model():
+    PRETRAINED_LM = "kkngan/bert-base-uncased-it-service-classification"
+    # model = AutoModelForSequenceClassification.from_pretrained(PRETRAINED_LM, num_labels=8)
+    # tokenizer = AutoTokenizer.from_pretrained(PRETRAINED_LM)
+    tokenizer = BertTokenizer.from_pretrained(PRETRAINED_LM, do_lower_case=True)
+    model = BertForSequenceClassification.from_pretrained(PRETRAINED_LM,
+                                                         num_labels=8)
+    return model, tokenizer
+
+
+def predict(text, model, tokenizer):
+    lookup_key ={0: 'Hardware',
+    1: 'Access',
+    2: 'Miscellaneous',
+    3: 'HR Support',
+    4: 'Purchase',
+    5: 'Administrative rights',
+    6: 'Storage',
+    7: 'Internal Project'}
+    inputs = tokenizer(text,
+                   padding = True,
+                   truncation = True,
+                   return_tensors='pt')
+    outputs = model(**inputs)
+    predicted_class_id = outputs.logits.argmax().item()
+    predicted_label = lookup_key.get(predicted_class_id)
+    probability = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().detach().numpy()
+    return predicted_label, predicted_class_id, probability
+
+
+def display_result(translate_text, prediction, predicted_class_id, probability):
+
+    category = ('Hardware',
+                'Access',
+                'Miscellaneous',
+                'HR Support',
+                'Purchase',
+                'Administrative rights',
+                'Storage',
+                'Internal Project')
+
+    # Show translated text and prediction
+    st.markdown('<font color="purple"><b>Text:</b></font>', unsafe_allow_html=True)
+    st.write(f'{translate_text}')
+    st.write(f'\n')
+    st.write(f'\n')
+    st.markdown('<font color="green"><b>Predicted Class:</b></font>', unsafe_allow_html=True)
+    st.write(f'{prediction}')
+    st.write(f'\n')
+    st.write(f'\n')
+
+    # Show Probability of each Service Category
+    
+    probability = np.array(probability[0])
+    df = pd.DataFrame({'Category': category, 'Probability (%)': probability * 100})
+    df['Probability (%)'] = df['Probability (%)'].apply(lambda x: round(x, 2))
+    base = alt.Chart(df).encode(
+                x='Probability (%)',
+                y=alt.Y('Category').sort('-x'),
+                tooltip=['Category',alt.Tooltip('Probability (%)', format=",.2f")],
+                text='Probability (%)'
+                ).properties(title="Probability of each Service Category")
+    chart = base.mark_bar() + base.mark_text(align='left', dx=2)
+    st.altair_chart(chart, use_container_width=True)
+
+
+def main():
+    # define parameters
+    image_path = 'front_page_image.jpg'
+    model_options = ["openai/whisper-base", "openai/whisper-large-v3"]
+    input_options = ["Start a recording", "Upload an audio", "Enter a transcript"]
+
+    # st.cache_resource.clear()
+    st.set_page_config(layout="wide", page_title="NLP IT Service Classification", page_icon="🤖",)
+    st.markdown('<b>🤖 Welcome to IT Service Classification Assistant!!! 🤖</b>', unsafe_allow_html=True)
+    st.write(f'\n')
+    st.write(f'\n')
+
+    with st.sidebar:
+        st.image(image_path , use_column_width=True)        
+        text_to_speech_model = st.selectbox("Pick select a speech to text model", model_options)
+        options = st.selectbox("Pick select an input method", input_options)
+
+        # start a recording
+        if options == input_options[0]:
+            audio = mic_recorder(key='my_recorder', callback=callback)
+        
+        # Upload an audio
+        elif options == input_options[1]:
+            audio = st.file_uploader("Please upload an audio", type=["wav", "mp3"])
+        
+        # Enter a transcript
+        else:
+            text = st.text_area("Please input the transcript (Only support English)")
+
+        button = st.button('Submit')
+
+    if button:        
+        with st.spinner(text="Loading... It may take a while if you are running the app for the first time."):
+
+            start_time = time.time()
+
+            # get inputs
+            if options == input_options[0]:
+                translate_text = translate(inputs=audio["bytes"], model=text_to_speech_model)
+
+            elif options == input_options[1]:
+                translate_text = translate(inputs=audio.getvalue(), model=text_to_speech_model)
+
+            else:
+                translate_text = text
+
+            model, tokenizer = load_classification_model()
+            prediction, predicted_class_id, probability = predict(text=translate_text, model=model, tokenizer=tokenizer)
+
+            end_time = time.time()
+
+        display_result(translate_text, prediction, predicted_class_id, probability)
+
+        st.write(f'\n')
+        st.write(f'\n')
+        st.markdown(f'*It took {(end_time-start_time):.2f} sec to process the input.', unsafe_allow_html=True)
+
+
+if __name__ == '__main__':
+    main()