add talking bot (#1216)

Author: Spycsh

workflows/chatbot/inference/backend/chat/README.md

@@ -73,4 +73,12 @@ or
 
 ```shell
 nohup bash run_itrex.sh &
+```
+
+# Install dependencies for TalkingBot
+
+```
+pip install speechbrain
+pip install soundfile
+pip install pydub
 ```

workflows/chatbot/inference/backend/chat/asr.py

@@ -0,0 +1,29 @@
+import torch
+import intel_extension_for_pytorch as ipex
+from transformers import WhisperForConditionalGeneration, WhisperProcessor
+from datasets import load_dataset, Audio, Dataset
+import time
+
+class AudioSpeechRecognition:
+ """Convert audio to text."""
+ def __init__(self):
+ self.device = "cpu"
+ self.model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-small").to(self.device)
+ self.processor = WhisperProcessor.from_pretrained("openai/whisper-small")
+ self.model.eval()
+ #self.model = ipex.optimize(self.model, dtype=torch.bfloat16)
+
+ def audio2text(self, audio_path):
+ """Convert audio to text
+
+ audio_path: the path to the input audio, e.g. ~/xxx.mp3
+ """
+ start = time.time()
+ audio_dataset = Dataset.from_dict({"audio": [audio_path]}).cast_column("audio", Audio(sampling_rate=16000))
+ waveform = audio_dataset[0]["audio"]['array']
+ inputs = self.processor.feature_extractor(waveform, return_tensors="pt", sampling_rate=16_000).input_features.to(self.device)
+ #with torch.cpu.amp.autocast():
+ predicted_ids = self.model.generate(inputs)
+ result = self.processor.tokenizer.batch_decode(predicted_ids, skip_special_tokens=True, normalize=True)[0]
+ print(f"generated text in {time.time() - start} seconds, and the result is: {result}")
+ return result

workflows/chatbot/inference/backend/chat/controller.py

@@ -13,8 +13,8 @@
 import threading
 import re
 
-from fastapi import FastAPI, Request
-from fastapi.responses import StreamingResponse
+from fastapi import FastAPI, Request, File, UploadFile, Form
+from fastapi.responses import StreamingResponse, FileResponse
 import numpy as np
 import requests
 import uvicorn
@@ -411,6 +411,27 @@ async def event_stream():
 
  return response
 
+@app.post("/v1/chat/talkingbot")
+async def handle_talkingbot(file: UploadFile = File(...), voice: str = Form(...)):
+ start = time.time()
+ file_name = file.filename
+ logger.info(f'Received file: {file_name}, and use voice: {voice}')
+ with open("tmp_audio_bytes", 'wb') as fout:
+ content = await file.read()
+ fout.write(content)
+ audio = AudioSegment.from_file("tmp_audio_bytes")
+ # bytes to mp3
+ audio.export(f"{file_name}", format="mp3")
+ worker_name = controller.get_worker_address("mpt-7b-chat")
+
+ try:
+ r = requests.post(worker_name + "/talkingbot", json={"file_name": file_name, "voice": voice}, timeout=20)
+ except requests.exceptions.RequestException as e:
+ logger.error(f"Talkingbot fails: {worker_name}, {e}")
+ return None
+ logger.info(f"E2E time: {time.time() - start}")
+ return FileResponse(r.content, media_type="video/mp4")
+
 
 if __name__ == "__main__":
  parser = argparse.ArgumentParser()

workflows/chatbot/inference/backend/chat/model_worker.py

@@ -14,14 +14,17 @@
 from fastapi import FastAPI, Request, BackgroundTasks
 from fastapi.responses import StreamingResponse
 import requests
-from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer
+from transformers import AutoTokenizer, AutoModelForCausalLM, LlamaTokenizer, GenerationConfig, StoppingCriteria, StoppingCriteriaList
 import torch
 import uvicorn
 
 from constants import WORKER_HEART_BEAT_INTERVAL
 from inference import load_model, generate_stream
 from utils import (build_logger, server_error_msg, pretty_print_semaphore)
 
+from asr import AudioSpeechRecognition
+from tts import TextToSpeech
+
 GB = 1 << 30
 
 worker_id = str(uuid.uuid4())[:6]
@@ -154,6 +157,88 @@ async def api_generate_stream(request: Request):
  background_tasks.add_task(release_model_semaphore)
  return StreamingResponse(generator, background=background_tasks, media_type="text/event-stream")
 
+class StopOnTokens(StoppingCriteria):
+ def __init__(self, min_length: int, start_length: int, stop_token_id: list[int]):
+ self.min_length = min_length
+ self.start_length = start_length
+ self.stop_token_id = stop_token_id
+
+ def __call__(
+ self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs
+ ) -> bool:
+ if scores is not None:
+ if len(scores) > self.min_length:
+ for stop_id in self.stop_token_id:
+ if input_ids[0][self.start_length - 1 + len(scores)] == stop_id:
+ return True
+ elif input_ids.shape[-1] - self.start_length > self.min_length:
+ for stop_id in self.stop_token_id:
+ if input_ids[0][input_ids.shape[-1] - 1] == stop_id:
+ return True
+ return False
+
+@app.post("/talkingbot", response_class=PlainTextResponse)
+async def talkingbot(request: Request):
+ params = await request.json()
+ saved_path = params["file_name"]
+ voice = params["voice"]
+ # audio -> text
+ logger.info("1: audio --> text")
+ text = asr.audio2text(saved_path)
+ logger.info(text)
+ prompt = """Have a conversation with a human. You must generate suitable response in short to the user input.\n### Input:\n{}\n### Response:""".format(text)
+ # text -> answer
+ logger.info("2: text --> answer")
+ worker.tokenizer.pad_token = worker.tokenizer.eos_token
+ stop_token_ids = [worker.model.model.generation_config.eos_token_id]
+ stop_token_ids.append(worker.tokenizer(".", return_tensors="pt").input_ids)
+ input_tokens = worker.tokenizer.batch_encode_plus([prompt], return_tensors="pt", padding=True)
+ input_token_len = input_tokens.input_ids.shape[-1]
+
+ stop = StopOnTokens(min_length=44, start_length=input_token_len, stop_token_id=stop_token_ids)
+ generation_config = GenerationConfig(
+ eos_token_id=0,
+ pad_token_id=0,
+ use_cache=True,
+ min_new_tokens=1,
+ max_new_tokens=64,
+ temperature=0.9,
+ top_p=0.9,
+ top_k=1,
+ repetition_penalty=1.1,
+ num_beams=1,
+ early_stopping=True,
+ ## use default decode mode
+ )
+ generation_kwargs = dict(
+ generation_config=generation_config, return_dict_in_generate=True
+ )
+ generation_kwargs["stopping_criteria"] = StoppingCriteriaList([stop])
+ with torch.no_grad():
+ with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True):
+ ## worker.model ==> ipexwrapper
+ output = worker.model.model.generate(**input_tokens, **generation_kwargs)
+ generated_texts = worker.tokenizer.decode(output.sequences[0], skip_special_tokens=True)
+ logger.info("raw generated texts", generated_texts)
+ if "### Response:" in generated_texts:
+ generated_texts = generated_texts.split("### Response:")[1].strip()
+ lines = generated_texts.split('\n')
+ result_lines = []
+ for line in lines:
+ if 'Input:' in line or '```python' in line:
+ break
+ result_lines.append(line)
+ generated_texts = '\n'.join(result_lines)
+ generated_texts = generated_texts.replace('#', '')
+ generated_texts = generated_texts.split('include <')[0]
+ # answer -> audio
+ # answer -> audio
+ logger.info("3: answer --> audio")
+ answer_speech_path = tts.text2speech(generated_texts, voice=voice)
+ logger.info("Done!!!")
+ logger.info(answer_speech_path)
+ return answer_speech_path
+
 
 @app.post("/worker_get_status")
 async def api_get_status(request: Request):
@@ -192,4 +277,6 @@ async def api_get_status(request: Request):
  args.load_8bit,
  args.itrex,
  args.ipex)
+ asr = AudioSpeechRecognition()
+ tts = TextToSpeech()
  uvicorn.run(app, host=args.host, port=args.port, log_level="info")

workflows/chatbot/inference/backend/chat/speaker_embeddings/spk_embed_default.pt

@@ -0,0 +1,109 @@
+from transformers import SpeechT5Processor, SpeechT5ForTextToSpeech, set_seed
+from datasets import load_dataset, Audio, Dataset, Features, ClassLabel
+import os
+import torch
+from speechbrain.pretrained import EncoderClassifier
+from dataclasses import dataclass
+from typing import Any, Dict, List, Union
+from transformers import SpeechT5HifiGan
+import soundfile as sf
+from datetime import datetime
+import intel_extension_for_pytorch as ipex
+import time
+import numpy as np
+from torch.utils.data import DataLoader
+
+class TextToSpeech:
+ """Convert text to speech with a driven speaker embedding
+
+ 1) Default voice (Original model + Proved good default speaker embedding from trained dataset)
+ 2) Finetuned voice (Fine-tuned offline model of specific person, such as Pat's voice + corresponding embedding)
+ 3) Customized voice (Original model + User's customized input voice embedding)
+ """
+ def __init__(self):
+ """Make sure your export LD_PRELOAD=<path to libiomp5.so and libtcmalloc> beforehand."""
+ # default setting
+ self.original_model = SpeechT5ForTextToSpeech.from_pretrained("microsoft/speecht5_tts")
+ #self.original_model = ipex.optimize(self.original_model, torch.bfloat16)
+ self.processor = SpeechT5Processor.from_pretrained("microsoft/speecht5_tts")
+ self.device = "cpu"
+ self.spk_model_name = "speechbrain/spkrec-xvect-voxceleb"
+ self.speaker_model = EncoderClassifier.from_hparams(
+ source=self.spk_model_name,
+ run_opts={"device": self.device},
+ savedir=os.path.join("/tmp", self.spk_model_name)
+ )
+ self.vocoder = SpeechT5HifiGan.from_pretrained("microsoft/speecht5_hifigan")
+ self.vocoder.eval()
+ self.default_speaker_embedding = torch.load('speaker_embeddings/spk_embed_default.pt') # load the default speaker embedding
+
+ # specific parameters for Pat Gelsinger
+ # preload the model in case of time-consuming runtime loading
+ self.pat_model = None
+ if os.path.exists("finetuned_model_1000_125_few_shot.pt"):
+ self.pat_model = torch.load("finetuned_model_1000_125_few_shot.pt", map_location=torch.device('cpu'))
+
+ # self.pat_model = ipex.optimize(self.pat_model, torch.bfloat16)
+ # self.speaker_embeddings = self.create_speaker_embedding(driven_audio_path)
+ self.pat_speaker_embeddings = None
+ if os.path.exists('speaker_embeddings/spk_embed_pat.pt'):
+ self.pat_speaker_embeddings = torch.load('speaker_embeddings/spk_embed_pat.pt')
+
+ # ipex IOMP hardware resources
+ self.cpu_pool = ipex.cpu.runtime.CPUPool([i for i in range(24)])
+
+ def create_speaker_embedding(self, driven_audio_path):
+ """Create the speaker's embedding
+
+ driven_audio_path: the driven audio of that speaker e.g. vgjwo-5bunm.mp3
+ """
+ audio_dataset = Dataset.from_dict({"audio": [driven_audio_path]}).cast_column("audio", Audio(sampling_rate=16000))
+ waveform = audio_dataset[0]["audio"]['array']
+ with torch.no_grad():
+ speaker_embeddings = self.speaker_model.encode_batch(torch.tensor(waveform))
+ speaker_embeddings = torch.nn.functional.normalize(speaker_embeddings, dim=2) # [1,1,512]
+ # speaker_embeddings = speaker_embeddings.squeeze().cpu().numpy()
+ speaker_embeddings = speaker_embeddings[0] # [1,512]
+ return speaker_embeddings.cpu()
+
+ def lookup_voice_embedding(self, voice):
+ if os.path.exists(f"speaker_embeddings/spk_embed_{voice}.pt") == False:
+ print("No customized speaker embedding is found! Use the default one")
+ return "speaker_embeddings/spk_embed_default.pt"
+ else:
+ return f"speaker_embeddings/spk_embed_{voice}.pt"
+
+ def text2speech(self, text, voice="default"):
+ """Text to speech.
+
+ text: the input text
+ voice: default/pat/huma/tom/eric...
+ """
+ start = time.time()
+ inputs = self.processor(text=text, return_tensors="pt")
+ model = self.original_model
+ speaker_embeddings = self.default_speaker_embedding
+
+ if voice == "pat":
+ if self.pat_model == None:
+ print("Finetuned model is not found! Use the default one")
+ else:
+ model = self.pat_model
+ if self.pat_speaker_embeddings == None:
+ print("Pat's speaker embedding is not found! Use the default one")
+ else:
+ speaker_embeddings = self.pat_speaker_embeddings
+ elif voice != "default":
+ speaker_embeddings = torch.load(self.lookup_voice_embedding(voice))
+
+ with torch.no_grad():
+ with ipex.cpu.runtime.pin(self.cpu_pool):
+ #with torch.cpu.amp.autocast(enabled=True, dtype=torch.bfloat16, cache_enabled=True):
+ spectrogram = model.generate_speech(inputs["input_ids"], speaker_embeddings)
+ speech = self.vocoder(spectrogram)
+ now = datetime.now()
+ time_stamp = now.strftime("%d_%m_%Y_%H_%M_%S")
+ output_video_path = f"output_{time_stamp}.wav"
+ print(f"text to speech in {time.time() - start} seconds, and dump the video at {output_video_path}")
+ sf.write(output_video_path, speech.cpu().numpy(), samplerate=16000)
+ return output_video_path