Spaces:
Running
on
T4
Running
on
T4
File size: 5,078 Bytes
2485dd8 |
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 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 |
# Creates a directory in which to look up available agents
from typing import List
from src.simuleval_transcoder import SimulevalTranscoder
import json
import logging
logger = logging.getLogger("socketio_server_pubsub")
# fmt: off
M4T_P0_LANGS = [
"eng",
"arb", "ben", "cat", "ces", "cmn", "cym", "dan",
"deu", "est", "fin", "fra", "hin", "ind", "ita",
"jpn", "kor", "mlt", "nld", "pes", "pol", "por",
"ron", "rus", "slk", "spa", "swe", "swh", "tel",
"tgl", "tha", "tur", "ukr", "urd", "uzn", "vie",
]
# fmt: on
class NoAvailableAgentException(Exception):
pass
class AgentWithInfo:
def __init__(
self,
agent,
name: str,
modalities: List[str],
target_langs: List[str],
# Supported dynamic params are defined in StreamingTypes.ts
dynamic_params: List[str] = [],
description="",
):
self.agent = agent
self.name = name
self.description = description
self.modalities = modalities
self.target_langs = target_langs
self.dynamic_params = dynamic_params
def get_capabilities_for_json(self):
return {
"name": self.name,
"description": self.description,
"modalities": self.modalities,
"targetLangs": self.target_langs,
"dynamicParams": self.dynamic_params,
}
@classmethod
def load_from_json(cls, config: str):
"""
Takes in JSON array of models to load in, e.g.
[{"name": "s2s_m4t_emma-unity2_multidomain_v0.1", "description": "M4T model that supports simultaneous S2S and S2T", "modalities": ["s2t", "s2s"], "targetLangs": ["en"]},
{"name": "s2s_m4t_expr-emma_v0.1", "description": "ES-EN expressive model that supports S2S and S2T", "modalities": ["s2t", "s2s"], "targetLangs": ["en"]}]
"""
configs = json.loads(config)
agents = []
for config in configs:
agent = SimulevalTranscoder.build_agent(config["name"])
agents.append(
AgentWithInfo(
agent=agent,
name=config["name"],
modalities=config["modalities"],
target_langs=config["targetLangs"],
)
)
return agents
class SimulevalAgentDirectory:
# Available models. These are the directories where the models can be found, and also serve as an ID for the model.
seamless_streaming_agent = "SeamlessStreaming"
def __init__(self):
self.agents = []
self.did_build_and_add_agents = False
def add_agent(self, agent: AgentWithInfo):
self.agents.append(agent)
def build_agent_if_available(self, model_id, config_name=None):
agent = None
try:
if config_name is not None:
agent = SimulevalTranscoder.build_agent(
model_id,
config_name=config_name,
)
else:
agent = SimulevalTranscoder.build_agent(
model_id,
)
except Exception as e:
logger.warning("Failed to build agent %s: %s" % (model_id, e))
raise e
return agent
def build_and_add_agents(self, models_override=None):
if self.did_build_and_add_agents:
return
if models_override is not None:
agent_infos = AgentWithInfo.load_from_json(models_override)
for agent_info in agent_infos:
self.add_agent(agent_info)
else:
s2s_m4t_expr_agent = self.build_agent_if_available(
SimulevalAgentDirectory.seamless_streaming_agent,
config_name="vad_s2st_sc_24khz_main.yaml",
)
if s2s_m4t_expr_agent:
self.add_agent(
AgentWithInfo(
agent=s2s_m4t_expr_agent,
name=SimulevalAgentDirectory.seamless_streaming_agent,
modalities=["s2t", "s2s"],
target_langs=M4T_P0_LANGS,
dynamic_params=["expressive"],
description="multilingual expressive model that supports S2S and S2T",
)
)
if len(self.agents) == 0:
logger.error(
"No agents were loaded. This likely means you are missing the actual model files specified in simuleval_agent_directory."
)
self.did_build_and_add_agents = True
def get_agent(self, name):
for agent in self.agents:
if agent.name == name:
return agent.agent
return None
def get_agent_or_throw(self, name):
agent = self.get_agent(name)
if agent is None:
raise NoAvailableAgentException("No agent found with name= %s" % (name))
return agent
def get_agents_capabilities_list_for_json(self):
return [agent.get_capabilities_for_json() for agent in self.agents]
|