mirror of
https://github.com/mudler/LocalAI.git
synced 2026-02-07 05:04:29 -05:00
148 lines
4.7 KiB
Python
148 lines
4.7 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
gRPC server of LocalAI for NVIDIA NEMO Toolkit ASR.
|
|
"""
|
|
from concurrent import futures
|
|
import time
|
|
import argparse
|
|
import signal
|
|
import sys
|
|
import os
|
|
import backend_pb2
|
|
import backend_pb2_grpc
|
|
import torch
|
|
import nemo.collections.asr as nemo_asr
|
|
|
|
import grpc
|
|
|
|
|
|
def is_float(s):
|
|
try:
|
|
float(s)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
def is_int(s):
|
|
try:
|
|
int(s)
|
|
return True
|
|
except ValueError:
|
|
return False
|
|
|
|
|
|
_ONE_DAY_IN_SECONDS = 60 * 60 * 24
|
|
MAX_WORKERS = int(os.environ.get('PYTHON_GRPC_MAX_WORKERS', '1'))
|
|
|
|
|
|
class BackendServicer(backend_pb2_grpc.BackendServicer):
|
|
def Health(self, request, context):
|
|
return backend_pb2.Reply(message=bytes("OK", 'utf-8'))
|
|
|
|
def LoadModel(self, request, context):
|
|
if torch.cuda.is_available():
|
|
device = "cuda"
|
|
else:
|
|
device = "cpu"
|
|
mps_available = hasattr(torch.backends, "mps") and torch.backends.mps.is_available()
|
|
if mps_available:
|
|
device = "mps"
|
|
if not torch.cuda.is_available() and request.CUDA:
|
|
return backend_pb2.Result(success=False, message="CUDA is not available")
|
|
|
|
self.device = device
|
|
self.options = {}
|
|
|
|
for opt in request.Options:
|
|
if ":" not in opt:
|
|
continue
|
|
key, value = opt.split(":", 1)
|
|
if is_float(value):
|
|
value = float(value)
|
|
elif is_int(value):
|
|
value = int(value)
|
|
elif value.lower() in ["true", "false"]:
|
|
value = value.lower() == "true"
|
|
self.options[key] = value
|
|
|
|
model_name = request.Model or "nvidia/parakeet-tdt-0.6b-v3"
|
|
|
|
try:
|
|
print(f"Loading NEMO ASR model from {model_name}", file=sys.stderr)
|
|
self.model = nemo_asr.models.ASRModel.from_pretrained(model_name=model_name)
|
|
print("NEMO ASR model loaded successfully", file=sys.stderr)
|
|
except Exception as err:
|
|
print(f"[ERROR] LoadModel failed: {err}", file=sys.stderr)
|
|
import traceback
|
|
traceback.print_exc(file=sys.stderr)
|
|
return backend_pb2.Result(success=False, message=str(err))
|
|
|
|
return backend_pb2.Result(message="Model loaded successfully", success=True)
|
|
|
|
def AudioTranscription(self, request, context):
|
|
result_segments = []
|
|
text = ""
|
|
try:
|
|
audio_path = request.dst
|
|
if not audio_path or not os.path.exists(audio_path):
|
|
print(f"Error: Audio file not found: {audio_path}", file=sys.stderr)
|
|
return backend_pb2.TranscriptResult(segments=[], text="")
|
|
|
|
# NEMO's transcribe method accepts a list of audio paths and returns a list of transcripts
|
|
results = self.model.transcribe([audio_path])
|
|
|
|
if not results or len(results) == 0:
|
|
return backend_pb2.TranscriptResult(segments=[], text="")
|
|
|
|
# Get the transcript text from the first result
|
|
text = results[0]
|
|
if text:
|
|
# Create a single segment with the full transcription
|
|
result_segments.append(backend_pb2.TranscriptSegment(
|
|
id=0, start=0, end=0, text=text
|
|
))
|
|
|
|
except Exception as err:
|
|
print(f"Error in AudioTranscription: {err}", file=sys.stderr)
|
|
import traceback
|
|
traceback.print_exc(file=sys.stderr)
|
|
return backend_pb2.TranscriptResult(segments=[], text="")
|
|
|
|
return backend_pb2.TranscriptResult(segments=result_segments, text=text)
|
|
|
|
|
|
def serve(address):
|
|
server = grpc.server(
|
|
futures.ThreadPoolExecutor(max_workers=MAX_WORKERS),
|
|
options=[
|
|
('grpc.max_message_length', 50 * 1024 * 1024),
|
|
('grpc.max_send_message_length', 50 * 1024 * 1024),
|
|
('grpc.max_receive_message_length', 50 * 1024 * 1024),
|
|
])
|
|
backend_pb2_grpc.add_BackendServicer_to_server(BackendServicer(), server)
|
|
server.add_insecure_port(address)
|
|
server.start()
|
|
print("Server started. Listening on: " + address, file=sys.stderr)
|
|
|
|
def signal_handler(sig, frame):
|
|
print("Received termination signal. Shutting down...")
|
|
server.stop(0)
|
|
sys.exit(0)
|
|
|
|
signal.signal(signal.SIGINT, signal_handler)
|
|
signal.signal(signal.SIGTERM, signal_handler)
|
|
|
|
try:
|
|
while True:
|
|
time.sleep(_ONE_DAY_IN_SECONDS)
|
|
except KeyboardInterrupt:
|
|
server.stop(0)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
parser = argparse.ArgumentParser(description="Run the gRPC server.")
|
|
parser.add_argument("--addr", default="localhost:50051", help="The address to bind the server to.")
|
|
args = parser.parse_args()
|
|
serve(args.addr)
|