Minor fixes to ds_tool and infer_tool

ultravox/tools/ds_tool.py

@@ -98,6 +98,7 @@ class DatasetToolArgs:
     num_workers: int = simple_parsing.field(default=16, alias="-w")
 
     upload_name: Optional[str] = simple_parsing.field(default=None, alias="-u")
+    upload_split: Optional[str] = simple_parsing.field(default=None)
     upload_branch: Optional[str] = simple_parsing.field(default="main", alias="-B")
     num_shards: Optional[int] = simple_parsing.field(default=None, alias="-N")
     private: bool = simple_parsing.field(default=False)
@@ -110,6 +111,11 @@ class DatasetToolArgs:
         positional=True,
     )
 
+    def __post_init__(self):
+        assert (
+            not self.upload_split or self.dataset_split
+        ), "Must specify dataset_split when using upload_split"
+
 
 def main(args: DatasetToolArgs):
     ds_name = args.dataset_name
@@ -132,6 +138,7 @@ def main(args: DatasetToolArgs):
         "token": token,
         "revision": args.upload_branch,
         "private": args.private,
+        "split": args.upload_split,
     }
     if args.num_shards is not None:
         hub_args["num_shards"] = {split: args.num_shards for split in data_dict.keys()}

ultravox/tools/infer_tool.py

@@ -78,6 +78,11 @@ class InferArgs:
     # JSON output
     json: bool = simple_parsing.field(default=False)
 
+    def __post_init__(self):
+        if self.prompt and self.prompt.startswith("@"):
+            with open(self.prompt[1:], "r") as f:
+                self.prompt = f.read()
+
 
 def run_tui(
     index: int,

ultravox/tools/tts.py

@@ -9,6 +9,8 @@
 import soundfile as sf
 
 RANDOM_VOICE_KEY = "random"
+REQUEST_TIMEOUT = 30
+NUM_RETRIES = 3
 
 
 def _make_ssml(voice: str, text: str):
@@ -23,14 +25,20 @@ def _make_ssml(voice: str, text: str):
 class Client(abc.ABC):
     def __init__(self, sample_rate: int = 16000):
         self._session = requests.Session()
+        retries = requests.adapters.Retry(total=NUM_RETRIES)
+        self._session.mount(
+            "https://", requests.adapters.HTTPAdapter(max_retries=retries)
+        )
         self._sample_rate = sample_rate
 
     @abc.abstractmethod
     def tts(self, text: str, voice: Optional[str] = None):
         raise NotImplementedError
 
     def _post(self, url: str, headers: Dict[str, str], json: Dict[str, Any]):
-        response = self._session.post(url, headers=headers, json=json)
+        response = self._session.post(
+            url, headers=headers, json=json, timeout=REQUEST_TIMEOUT
+        )
         response.raise_for_status()
         return response
 
@@ -132,7 +140,6 @@ def tts(self, text: str, voice: Optional[str] = None):
             i = np.random.randint(len(self.ALL_VOICES)) + os.getpid()
             voice = self.ALL_VOICES[i % len(self.ALL_VOICES)]
         url = f"https://api.elevenlabs.io/v1/text-to-speech/{voice}/stream?output_format=pcm_16000"
-        print("url", url)
         headers = {"xi-api-key": os.environ["ELEVEN_API_KEY"]}
         body = {
             "text": text,