Update app.py

app.py

@@ -3,20 +3,24 @@ from streamlit_mic_recorder import mic_recorder
 from transformers import pipeline
 import torch
 from transformers import BertTokenizer, BertForSequenceClassification, AutoModelForSequenceClassification, AutoTokenizer
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
 import numpy as np
 import pandas as pd
+import time
+
 
 def callback():
     if st.session_state.my_recorder_output:
         audio_bytes = st.session_state.my_recorder_output['bytes']
         st.audio(audio_bytes)
 
-def transcribe_and_translate(upload):
-    # pipe = pipeline("automatic-speech-recognition", model="openai/whisper-large")
-    pipe = pipeline("automatic-speech-recognition", model="openai/whisper-large-v3")
-    transcribe_result = pipe(upload, generate_kwargs={'task': 'transcribe'})
-    translate_result = pipe(upload, generate_kwargs={'task': 'translate'})
-    return transcribe_result['text'], translate_result['text']
+
+def translate(inputs, model="openai/whisper-medium"):
+    pipe = pipeline("automatic-speech-recognition", model=model)
+    # transcribe_result = pipe(upload, generate_kwargs={'task': 'transcribe'})
+    translate_result = pipe(inputs, generate_kwargs={'task': 'translate'})
+    return translate_result['text']
+
 
 def encode_depracated(docs, tokenizer):
     '''
@@ -29,16 +33,16 @@ def encode_depracated(docs, tokenizer):
     return input_ids, attention_masks
 
 
-def load_model():
-    CUSTOMMODEL_PATH = "./bert-itserviceclassification"
-    PRETRAINED_LM = "bert-base-uncased"
-    tokenizer = BertTokenizer.from_pretrained(PRETRAINED_LM, do_lower_case=True)
-    model = BertForSequenceClassification.from_pretrained(PRETRAINED_LM,
-                                                        num_labels=8,
-                                                        output_attentions=False,
-                                                        output_hidden_states=False)
-    model.load_state_dict(torch.load(CUSTOMMODEL_PATH, map_location ='cpu'))
-    return model, tokenizer
+# def load_model_deprecated():
+#     CUSTOMMODEL_PATH = "./bert-itserviceclassification"
+#     PRETRAINED_LM = "bert-base-uncased"
+#     tokenizer = BertTokenizer.from_pretrained(PRETRAINED_LM, do_lower_case=True)
+#     model = BertForSequenceClassification.from_pretrained(PRETRAINED_LM,
+#                                                         num_labels=8,
+#                                                         output_attentions=False,
+#                                                         output_hidden_states=False)
+#     model.load_state_dict(torch.load(CUSTOMMODEL_PATH, map_location ='cpu'))
+#     return model, tokenizer
 
 
 def load_model():
@@ -67,18 +71,19 @@ def predict(text, model, tokenizer):
     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, probability
+    confidence = torch.nn.functional.softmax(outputs.logits, dim=-1).cpu().detach().numpy()
+    return predicted_label, confidence
 
 
 def main():
-
     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('front_page_image.jpg' , use_column_width=True)
+        text_to_speech_model = st.selectbox("Pick select a text_to_speech_model", ["openai/whisper-base", "openai/whisper-medium", "openai/whisper-large-v3"])
         options = st.selectbox("Pick select an input method", ["Start a recording", "Upload an audio", "Enter a transcript"])
         if options == "Start a recording":
             audio = mic_recorder(key='my_recorder', callback=callback)
@@ -89,32 +94,46 @@ def main():
         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()
             model, tokenizer = load_model()
             if options == "Start a recording":
-                transcibe_text, translate_text = transcribe_and_translate(upload=audio["bytes"])
-                prediction, probability = predict(text=translate_text, model=model, tokenizer=tokenizer)
+                # transcibe_text, translate_text = transcribe_and_translate(upload=audio["bytes"])
+                translate_text = translate(inputs=audio["bytes"], model=text_to_speech_model)
+                prediction, confidence = predict(text=translate_text, model=model, tokenizer=tokenizer)
             elif options == "Upload an audio":
-                transcibe_text, translate_text = transcribe_and_translate(upload=audio.getvalue())
-                prediction, probability = predict(text=translate_text, model=model, tokenizer=tokenizer)
+                # transcibe_text, translate_text = transcribe_and_translate(upload=audio.getvalue())
+                translate_text = translate(inputs=audio.getvalue(), model=text_to_speech_model)
+                prediction, confidence = predict(text=translate_text, model=model, tokenizer=tokenizer)
             else:
-                transcibe_text = text
-                prediction, probability = predict(text=text, model=model, tokenizer=tokenizer)
-        st.markdown('<font color="blue"><b>Transcript:</b></font>', unsafe_allow_html=True)
-        st.write(f'{transcibe_text}')
+                translate_text = text
+                prediction, confidence = predict(text=text, model=model, tokenizer=tokenizer)
+            end_time = time.time()
+        # st.markdown('<font color="blue"><b>Transcript:</b></font>', unsafe_allow_html=True)
+        # st.write(f'{transcibe_text}')
+        # st.write(f'\n')
+        # if options != "Enter a transcript":
+        st.markdown('<font color="purple"><b>(Translated) Text:</b></font>', unsafe_allow_html=True)
+        st.write(f'{translate_text}')
+        st.write(f'\n')
         st.write(f'\n')
-        if options != "Enter a transcript":
-            st.markdown('<font color="red"><b>Translation:</b></font>', unsafe_allow_html=True)
-            st.write(f'{translate_text}')
-            st.write(f'\n')
         st.markdown('<font color="green"><b>Predicted Class:</b></font>', unsafe_allow_html=True)
         st.write(f'{prediction}')
 
-        # Convert probability to bar
+        # Convert confidence to bar cart
+        st.write(f'\n')
+        st.write(f'\n')
+        category = ('Hardware', 'Access', 'Miscellaneous', 'HR Support', 'Purchase', 'Administrative rights', 'Storage', 'Internal Project')
+        confidence = np.array(confidence[0])
+        df = pd.DataFrame({'Category': category, 'Confidence (%)': confidence * 100})
+        df['Confidence (%)'] = df['Confidence (%)'].apply(lambda x: round(x, 2))
+        st.bar_chart(data=df, x='Category', y='Confidence (%)')
+        # df = df.sort_values(by='Confidence (%)', ascending=False).reset_index(drop=True)
+        # st.write(df)
+        st.write(f'\n')
         st.write(f'\n')
-        objects = ('Hardware', 'Access', 'Miscellaneous', 'HR Support', 'Purchase', 'Administrative rights', 'Storage', 'Internal Project')
-        df = pd.DataFrame({'Categories': objects, 'Probability': np.around(probability[0])})
-        st.bar_chart(data=df, x='Categories', y='Probability')
+        st.markdown(f'*It took {(end_time-start_time):.2f} sec to process the input', unsafe_allow_html=True)
 
 if __name__ == '__main__':
     main()