Update functions.py

functions.py

@@ -6,7 +6,7 @@ import plotly_express as px
 import nltk
 import plotly.graph_objects as go
 from optimum.onnxruntime import ORTModelForSequenceClassification
-from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification,WhisperProcessor, WhisperForConditionalGeneration
+from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
 from sentence_transformers import SentenceTransformer, CrossEncoder, util
 import streamlit as st
 import en_core_web_lg
@@ -29,8 +29,8 @@ margin-bottom: 2.5rem">{}</div> """
 
 @st.experimental_singleton(suppress_st_warning=True)
 def load_models():
-    #asr_model = whisper.load_model("small")
-    asr_pipe = pipeline("automatic-speech-recognition",model = "openai/whisper-small")
+    asr_model = whisper.load_model("base")
+    #asr_pipe = pipeline("automatic-speech-recognition",model = "openai/whisper-small")
     q_model = ORTModelForSequenceClassification.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
     ner_model = AutoModelForTokenClassification.from_pretrained("xlm-roberta-large-finetuned-conll03-english")
     q_tokenizer = AutoTokenizer.from_pretrained("nickmuchi/quantized-optimum-finbert-tone")
@@ -41,7 +41,7 @@ def load_models():
     sbert = SentenceTransformer("all-mpnet-base-v2")
     cross_encoder = CrossEncoder('cross-encoder/ms-marco-MiniLM-L-12-v2')
     
-    return asr_pipe, sent_pipe, sum_pipe, ner_pipe, sbert, cross_encoder
+    return asr_model, sent_pipe, sum_pipe, ner_pipe, sbert, cross_encoder
 
 @st.experimental_singleton(suppress_st_warning=True)
 def get_spacy():
@@ -57,12 +57,12 @@ def inference(link, upload):
       yt = YouTube(link)
       title = yt.title
       path = yt.streams.filter(only_audio=True)[0].download(filename="audio.mp4")
-      results = asr_pipe("audio.mp4")
+      results = asr_model.transcribe(path)
       
       return results['text'], yt.title
       
     elif upload:
-      results = asr_pipe(upload)
+      results = asr_model.trasncribe(upload)
       
       return results['text'], "Transcribed Earnings Audio"
 
@@ -306,4 +306,4 @@ def fin_ext(text):
     return make_spans(text,results)
     
 nlp = get_spacy()    
-asr_pipe, sent_pipe, sum_pipe, ner_pipe, sbert, cross_encoder  = load_models()
+asr_model, sent_pipe, sum_pipe, ner_pipe, sbert, cross_encoder  = load_models()