app.py

import streamlit as st
from faster_whisper import WhisperModel
import datetime
import subprocess
from pathlib import Path
import pandas as pd
import re
import time
import os
import numpy as np
from sklearn.cluster import AgglomerativeClustering
from sklearn.metrics import silhouette_score

import torch
from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
from pyannote.audio import Audio
from pyannote.core import Segment


import wave
import contextlib
from transformers import pipeline
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer
import onnxruntime
import numpy as np
import librosa

whisper_models = ["tiny", "base", "small", "medium", "large-v1", "large-v2"]
source_languages = {"en": "English"}

MODEL_NAME = "vumichien/whisper-medium-jp"
lang = "en"

device = 0 if torch.cuda.is_available() else "cpu"

embedding_model = PretrainedSpeakerEmbedding(
    "speechbrain/spkrec-ecapa-voxceleb",
    device=torch.device("cuda"))


#LLAMA prep
# from huggingface_hub import login

# login("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")

# import torch
# import transformers
# from transformers import AutoTokenizer, AutoModelForCausalLM
# from langchain import HuggingFacePipeline
# from langchain import PromptTemplate,  LLMChain


# tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-13b-chat-hf",
#                                           use_auth_token=True,)

# model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-2-13b-chat-hf",
#                                              device_map='auto',
#                                              torch_dtype=torch.float16,
#                                              use_auth_token=True,
#                                             #  load_in_8bit=True,
#                                             #  load_in_4bit=True
#                                              )
# # Use a pipeline for later
# from transformers import pipeline

# pipe = pipeline("text-generation",
#                 model=model,
#                 tokenizer= tokenizer,
#                 torch_dtype=torch.bfloat16,
#                 device_map="auto",
#                 max_new_tokens = 512,
#                 do_sample=True,
#                 top_k=30,
#                 num_return_sequences=1,
#                 eos_token_id=tokenizer.eos_token_id
#                 )


# import json
# import textwrap

# B_INST, E_INST = "[INST]", "[/INST]"
# B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
# DEFAULT_SYSTEM_PROMPT = """\
# You are a helpful, respectful and honest assistant. Always answer as helpfully as possible, while being safe. Your answers should not include any harmful, unethical, racist, sexist, toxic, dangerous, or illegal content. Please ensure that your responses are socially unbiased and positive in nature.

# If a question does not make any sense, or is not factually coherent, explain why instead of answering something not correct. If you don't know the answer to a question, please don't share false information."""



# def get_prompt(instruction, new_system_prompt=DEFAULT_SYSTEM_PROMPT ):
#     SYSTEM_PROMPT = B_SYS + new_system_prompt + E_SYS
#     prompt_template =  B_INST + SYSTEM_PROMPT + instruction + E_INST
#     return prompt_template

# def cut_off_text(text, prompt):
#     cutoff_phrase = prompt
#     index = text.find(cutoff_phrase)
#     if index != -1:
#         return text[:index]
#     else:
#         return text

# def remove_substring(string, substring):
#     return string.replace(substring, "")



# def generate(text):
#     prompt = get_prompt(text)
#     with torch.autocast('cuda', dtype=torch.bfloat16):
#         inputs = tokenizer(prompt, return_tensors="pt").to('cuda')
#         outputs = model.generate(**inputs,
#                                  max_new_tokens=512,
#                                  eos_token_id=tokenizer.eos_token_id,
#                                  pad_token_id=tokenizer.eos_token_id,
#                                  )
#         final_outputs = tokenizer.batch_decode(outputs, skip_special_tokens=True)[0]
#         final_outputs = cut_off_text(final_outputs, '</s>')
#         final_outputs = remove_substring(final_outputs, prompt)

#     return final_outputs#, outputs

# def parse_text(text):
#         wrapped_text = textwrap.fill(text, width=100)
#         print(wrapped_text +'\n\n')
#         # return assistant_text

# llm = HuggingFacePipeline(pipeline = pipe, model_kwargs = {'temperature':0})

def segment_embedding(segment, duration, audio_file):
  audio = Audio()
  start = segment["start"]
  # Whisper overshoots the end timestamp in the last segment
  end = min(duration, segment["end"])
  clip = Segment(start, end)
  waveform, sample_rate = audio.crop(audio_file, clip)
  return embedding_model(waveform[None])

def fast_whisper(audio_file, model):
  # Transcribe audio
  options = dict(language=lang, beam_size=5, best_of=5)
  transcribe_options = dict(task="transcribe", **options)
  segments_raw, info = model.transcribe(audio_file, **transcribe_options)

  # Convert back to original openai format
  segments = []
  i = 0
  for segment_chunk in segments_raw:
      chunk = {}
      chunk["start"] = segment_chunk.start
      chunk["end"] = segment_chunk.end
      chunk["text"] = segment_chunk.text
      segments.append(chunk)
      i += 1
  print("transcribe audio done with fast whisper")

  return segments

def get_embeddings(segments, duration, audio_file):
  embeddings = np.zeros(shape=(len(segments), 192))
  for i, segment in enumerate(segments):
      embeddings[i] = segment_embedding(segment, duration, audio_file)
  embeddings = np.nan_to_num(embeddings)

  print("Got embeddings for segments")
  return embeddings

def get_n_speakers(embeddings, num_speakers):
  if num_speakers == 0:
    # Find the best number of speakers
    score_num_speakers = {}

    for num_speakers in range(2, 10+1):
        clustering = AgglomerativeClustering(num_speakers).fit(embeddings)
        score = silhouette_score(embeddings, clustering.labels_, metric='euclidean')
        score_num_speakers[num_speakers] = score
    best_num_speaker = max(score_num_speakers, key=lambda x:score_num_speakers[x])
    print(f"The best number of speakers: {best_num_speaker} with {score_num_speakers[best_num_speaker]} score")
  else:
      best_num_speaker = num_speakers

  print(f"best num speakers is {best_num_speaker}")

  return best_num_speaker

def assign_speaker(best_num_speaker, embeddings, segments):
  # Assign speaker label
  clustering = AgglomerativeClustering(best_num_speaker).fit(embeddings)
  labels = clustering.labels_
  for i in range(len(segments)):
      segments[i]["speaker"] = 'SPEAKER ' + str(labels[i] + 1)

  print(f"I know who said what now")
  return segments

def convert_time(secs):
    return datetime.timedelta(seconds=round(secs))

def segments2df(segments):
  # Make output
  objects = {
      'Start' : [],
      'End': [],
      'Speaker': [],
      'Text': []
  }
  text = ''
  for (i, segment) in enumerate(segments):
      if i == 0 or segments[i - 1]["speaker"] != segment["speaker"]:
          objects['Start'].append(str(convert_time(segment["start"])))
          objects['Speaker'].append(segment["speaker"])
          if i != 0:
              objects['End'].append(str(convert_time(segments[i - 1]["end"])))
              objects['Text'].append(text)
              text = ''
      text += segment["text"] + ' '
  objects['End'].append(str(convert_time(segments[i - 1]["end"])))
  objects['Text'].append(text)

  df_results = pd.DataFrame(objects)

  return df_results


def speech_to_text(audio_file, whisper_model, num_speakers=0):
  model = WhisperModel(whisper_model, compute_type="int8")
  time_start = time.time()

  if(audio_file == None): raise ValueError("Error no audio_file")

  model = WhisperModel(whisper_model, compute_type="int8")
  y, sr = librosa.load(audio_file)
  duration = len(y)/sr
  segments = fast_whisper(audio_file, model)
  embeddings = get_embeddings(segments, duration, audio_file)
  best_num_speaker = get_n_speakers(embeddings, num_speakers)
  segments = assign_speaker(best_num_speaker, embeddings, segments)
  diary = segments2df(segments)

  return diary

onnx_path = hf_hub_download(repo_id='UKP-SQuARE/roberta-base-pf-boolq-onnx', filename='model.onnx') # or model_quant.onnx for quantization
onnx_model = onnxruntime.InferenceSession(onnx_path, providers=['CPUExecutionProvider'])

question = 'Can she answer'
tokenizer = AutoTokenizer.from_pretrained('UKP-SQuARE/roberta-base-pf-boolq-onnx')

def answer(context, question):
    # inputs = tokenizer(question, context, padding=True, truncation=True, return_tensors='np')
    # inputs = {key: np.array(inputs[key], dtype=np.int64) for key in inputs}
    # outputs = onnx_model.run(input_feed=dict(inputs), output_names=None)

    # instruction = f"conversation: '''{context}'''"+"\n based on the provided conversation in triple quotes answer next question.\n Question: {text}"

    # system_prompt = "You are an expert and answer any question based on conversation. You analys the conversation in light of the question then you answer with yes, no or not clear only. You only output one or two words"
    
    # template = get_prompt(instruction, system_prompt)
    # print(template)
    
    # prompt = PromptTemplate(template=template, input_variables=["text"])
    # llm_chain = LLMChain(prompt=prompt, llm=llm)
    # output = llm_chain.run(question)

    # return parse_text(output)
    return "please use the other app"


uploaded_file = st.sidebar.file_uploader("Choose a file")
num_speakers = st.sidebar.slider("num speakers (0 means auto detect)", 0, 10, 0)
diary = None
question = None
if uploaded_file is not None:
  filename = uploaded_file.name
  with open(filename, "wb") as f:
    f.write(uploaded_file.getbuffer())
  # st.write(os.listdir("./"))

if st.sidebar.checkbox('Get conversation'):
  torch.cuda.empty_cache()
  whisper_model = "base"
  diary = speech_to_text(filename, whisper_model, num_speakers=num_speakers)

  st.dataframe(diary.style.highlight_max(axis=0))


  question = st.sidebar.text_input('Question', 'Can she answer')
  if st.sidebar.button('Answer'):
    diary["text_all"] = diary["Speaker"] + ": "+ diary["Text"]
    context = " \n ".join(diary["text_all"].to_list())
    outputs = answer(context, question)

    outputs = outputs[0][0]
    if outputs[0]>outputs[1]: st.sidebar.write("Answer is Yes")
    if outputs[0]<outputs[1]: st.sidebar.write("Answer is No")