added ner

app.py

@@ -2,36 +2,43 @@ import streamlit as st
 from st_audiorec import st_audiorec
 
 from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
-from datasets import load_dataset
+#from datasets import load_dataset
 import torch
+from gliner import GLiNER
 
-pipe = None
-audio_sample: bytes = None
-audio_transcription: str = None
+from resources import Lead_Labels, entity_labels, set_start, audit_elapsedtime
 
 def main ():
     
-    print("Run init model")
-    pipe = init_model()
-    # x = st.slider('Select a value')
-    # st.write(x, 'squared is', x * x)
+    
+    rec = init_model_trans()
+    ner = init_model_ner() #async
+    
+    labels = entity_labels
+
+    text = "I have a proposal from cgd where they want one outsystems junior developers and one senior for an estimate of three hundred euros a day, for six months."
+    print(f"get entities from sample text: {text}")
+    get_entity_labels(model=ner, text=text, labels=labels)
 
     print("Render UI")
     wav_audio_data = st_audiorec()
 
     if wav_audio_data is not None:
         print("Loading data...")
+    
+    if wav_audio_data is not None and rec is not None:
+        print("Loading data...")
         st.audio(wav_audio_data, format='audio/wav')
-        transcribe(wav_audio_data, pipe)
-        
+        text = transcribe(wav_audio_data, rec)
+        if text is not None:    
+            get_entity_labels(labels=labels, model=ner, text=text)
 
-    
-    # dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
-    # sample = dataset[0]["audio"]
-    
-    # audio_file_path = "data/audio1.wav"
 
-def init_model ():
+def init_model_trans ():
+    print("Initiating transcription model...")
+    func_name = "init_model_trans"
+    start = set_start()
+
     device = "cuda:0" if torch.cuda.is_available() else "cpu"
     torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
 
@@ -57,8 +64,16 @@ def init_model ():
         device=device,
     )
     print(f'Init model successful: {model}' )
+    audit_elapsedtime(function=func_name, start=start)
     return pipe
-    
+
+async def init_model_ner():
+    print("Initiating NER model...")
+    start = set_start()
+    model = GLiNER.from_pretrained("urchade/gliner_multi")
+    audit_elapsedtime(function="init_model_ner", start=start)
+    return model
+
 def transcribe (audio_sample: bytes, pipe) -> str:
 
     # dataset = load_dataset("distil-whisper/librispeech_long", "clean", split="validation")
@@ -66,7 +81,16 @@ def transcribe (audio_sample: bytes, pipe) -> str:
     result = pipe(audio_sample)
     print(result)
     
-    st.write('Result', result["text"])
+    st.write('trancription: ', result["text"])
+    return result["text"]
+
+def get_entity_labels(model: GLiNER, text: str, labels: list): #-> Lead_labels:
+    entities = model.predict_entities(text, labels)
+
+    for entity in entities:
+        print(entity["text"], "=>", entity["label"])
+    st.write('Entities: ', entities)
+    # return Lead_Labels()
 
 if __name__ == "__main__":
     main()