Spaces:
Paused
Paused
File size: 3,698 Bytes
7cc82ad f1e930a 7cc82ad f1e930a e3f7cd8 f2683ae f1e930a 5d7959c f1e930a 5d7959c f1e930a 5d7959c f1e930a 7cc82ad f1e930a 5d7959c f1e930a e5e84e9 7cc82ad f2683ae 7cc82ad f1e930a 7cc82ad e3f7cd8 e5e84e9 f1e930a 5d7959c f1e930a f2683ae f1e930a e5e84e9 7cc82ad f2683ae 16388cf f2683ae e5e84e9 7cc82ad 16388cf |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 |
import multiprocessing
import argparse
import threading
import ssl
import time
import sys
import functools
from multiprocessing import Process, Manager, Value, Queue
from whisper_live.trt_server import TranscriptionServer
from llm_service import TensorRTLLMEngine
from tts_service import WhisperSpeechTTS
def parse_arguments():
parser = argparse.ArgumentParser()
parser.add_argument('--whisper_tensorrt_path',
type=str,
default="/root/TensorRT-LLM/examples/whisper/whisper_small_en",
help='Whisper TensorRT model path')
parser.add_argument('--mistral',
action="store_true",
help='Mistral')
parser.add_argument('--mistral_tensorrt_path',
type=str,
default=None,
help='Mistral TensorRT model path')
parser.add_argument('--mistral_tokenizer_path',
type=str,
default="teknium/OpenHermes-2.5-Mistral-7B",
help='Mistral TensorRT model path')
parser.add_argument('--phi',
action="store_true",
help='Phi')
parser.add_argument('--phi_tensorrt_path',
type=str,
default="/root/TensorRT-LLM/examples/phi/phi_engine",
help='Phi TensorRT model path')
parser.add_argument('--phi_tokenizer_path',
type=str,
default="/root/TensorRT-LLM/examples/phi/phi-2",
help='Phi Tokenizer path')
return parser.parse_args()
if __name__ == "__main__":
args = parse_arguments()
if not args.whisper_tensorrt_path:
raise ValueError("Please provide whisper_tensorrt_path to run the pipeline.")
import sys
sys.exit(0)
if args.mistral:
if not args.mistral_tensorrt_path or not args.mistral_tokenizer_path:
raise ValueError("Please provide mistral_tensorrt_path and mistral_tokenizer_path to run the pipeline.")
import sys
sys.exit(0)
if args.phi:
if not args.phi_tensorrt_path or not args.phi_tokenizer_path:
raise ValueError("Please provide phi_tensorrt_path and phi_tokenizer_path to run the pipeline.")
import sys
sys.exit(0)
multiprocessing.set_start_method('spawn')
lock = multiprocessing.Lock()
manager = Manager()
shared_output = manager.list()
transcription_queue = Queue()
llm_queue = Queue()
audio_queue = Queue()
whisper_server = TranscriptionServer()
whisper_process = multiprocessing.Process(
target=whisper_server.run,
args=(
"0.0.0.0",
6006,
transcription_queue,
llm_queue,
args.whisper_tensorrt_path
)
)
whisper_process.start()
llm_provider = TensorRTLLMEngine()
# llm_provider = MistralTensorRTLLMProvider()
llm_process = multiprocessing.Process(
target=llm_provider.run,
args=(
# args.mistral_tensorrt_path,
# args.mistral_tokenizer_path,
args.phi_tensorrt_path,
args.phi_tokenizer_path,
transcription_queue,
llm_queue,
audio_queue,
)
)
llm_process.start()
# audio process
tts_runner = WhisperSpeechTTS()
tts_process = multiprocessing.Process(target=tts_runner.run, args=("0.0.0.0", 8888, audio_queue))
tts_process.start()
llm_process.join()
whisper_process.join()
tts_process.join()
|