Update gradio demo to support text/voice conversation

.gitignore

@@ -125,6 +125,7 @@ ipython_config.py
 #   commonly ignored for libraries.
 #   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
 #poetry.lock
+poetry.toml
 
 # pdm
 #   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.

Justfile

@@ -8,6 +8,7 @@ default: format check test
 
 install:
     pip install poetry==1.7.1
+    poetry config virtualenvs.in-project true --local
     poetry install
 
 format:

pyproject.toml

@@ -9,13 +9,13 @@ readme = "README.md"
 
 [tool.poetry.dependencies]
 python = "^3.11"
-torch = "2.2.2"
+torch = "2.4"
 transformers = {version = ">=4.43.1", extras = ["torch"]}
 bitsandbytes = "~0.42.0"
 peft = "~0.11.1"
 simple-parsing = "~0.1.5"
 librosa = "~0.10.2.post1"
-requests = "~2.26.0"
+requests = "~2.31.0"
 datasets = "~2.19.1"
 mosaicml-streaming = "~0.7.6"
 nltk = "~3.8.1"
@@ -39,8 +39,8 @@ fsspec = "~2024.3.1"
 gcsfs = "~2024.3.1"
 sounddevice = "~0.4.7"
 mosaicml-cli = "~0.6.31"
-gradio-client = "~1.0.1"
-gradio = "~3.40.1"
+gradio-client = "~0.16.1"
+gradio = "~4.29.0"
 gpustat = "~1.1.1"
 types-requests = "~2.26.0"
 types-pyyaml = "^6.0.12.20240724"

ultravox/data/datasets.py

@@ -191,6 +191,9 @@ def __post_init__(self):
             ), f"Unexpected audio dtype: {self.audio.dtype}"
             assert self.audio.ndim == 1, f"Unexpected audio shape: {self.audio.shape}"
 
+    def add_past_messages(self, past_messages: List[Dict[str, str]]):
+        self.messages = past_messages + self.messages
+
     messages: List[Dict[str, str]]
     """List of messages, each with a "role" and "content" field."""
     audio: Optional[np.typing.NDArray[np.float32]] = None

ultravox/inference/base.py

@@ -1,6 +1,8 @@
 import abc
 import dataclasses
-from typing import Generator, Optional
+from typing import Dict, Generator, List, Optional, Tuple, Union
+
+import transformers
 
 from ultravox.data import datasets
 
@@ -10,6 +12,8 @@ class VoiceOutput:
     text: str
     input_tokens: int
     output_tokens: int
+    audio_token_len: int = 0
+    past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None
 
 
 class InferenceMessage:
@@ -37,6 +41,7 @@ def infer(
         sample: datasets.VoiceSample,
         max_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
     ) -> VoiceOutput:
         pass
 
@@ -45,8 +50,51 @@ def infer_stream(
         sample: datasets.VoiceSample,
         max_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
     ) -> InferenceGenerator:
         """Streaming polyfill, if not supported directly in derived classes."""
-        output = self.infer(sample, max_tokens, temperature)
+        output = self.infer(sample, max_tokens, temperature, past_key_values)
         yield InferenceChunk(output.text)
         yield InferenceStats(output.input_tokens, output.output_tokens)
+
+
+class History:
+    def __init__(self, audio_token_replacement: str = "<|eot_token|>"):
+        self.audio_token_replacement: str = audio_token_replacement
+        self.audio_placeholder = "<|audio|>"
+        self.messages: List[Dict[str, str]] = []
+        self.key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None
+
+    def add_message(self, message: Dict[str, str], audio_token_len: int):
+        message = message.copy()
+        content = message["content"]
+        if audio_token_len > 0:
+            if content.count(self.audio_placeholder) != 1:
+                raise ValueError(
+                    f"Expected 1 audio placeholder, found {content.count(self.audio_placeholder)}"
+                )
+            message["content"] = content.replace(
+                self.audio_placeholder, self.audio_token_replacement * audio_token_len
+            )
+
+        if self.messages:
+            self.messages.append(message)
+        else:
+            self.messages = [message]
+
+    def update_key_values(
+        self, key_values: Union[Tuple, transformers.cache_utils.Cache]
+    ):
+        self.key_values = key_values
+
+    @property
+    def past_messages(self) -> List[Dict[str, str]]:
+        return self.messages
+
+    @property
+    def past_key_values(self) -> Optional[Union[Tuple, transformers.cache_utils.Cache]]:
+        return self.key_values
+
+    def reset(self):
+        self.messages = []
+        self.key_values = None

ultravox/inference/infer.py

@@ -1,5 +1,5 @@
 import threading
-from typing import Optional
+from typing import Optional, Tuple, Union
 
 import librosa
 import numpy as np
@@ -11,7 +11,7 @@
 from ultravox.model import ultravox_processing
 
 SAMPLE_RATE = 16000
-MAX_TOKENS = 1024
+MAX_NEW_TOKENS = 1024
 # Without this penalty, the model tends to repeat itself.
 REPETITION_PENALTY = 1.1
 
@@ -33,22 +33,33 @@ def __init__(
     def infer(
         self,
         sample: datasets.VoiceSample,
-        max_tokens: Optional[int] = None,
+        max_new_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        num_beams: int = 1,
     ) -> base.VoiceOutput:
         inputs = self._dataproc(sample)
         input_len = inputs["input_ids"].shape[1]
-        output = self._generate(inputs, max_tokens, temperature)
-        output_tokens = output[0][input_len:]
+        output = self._generate(
+            inputs, max_new_tokens, temperature, past_key_values, num_beams
+        )
+        output_tokens = output.sequences[0][input_len:]
         output_text = self.tokenizer.decode(output_tokens, skip_special_tokens=True)
         output_len = len(output_tokens)
-        return base.VoiceOutput(output_text, input_len, output_len)
+        audio_token_len = 0
+        if "audio_token_len" in inputs:
+            audio_token_len = inputs["audio_token_len"][0]
+        return base.VoiceOutput(
+            output_text, input_len, output_len, audio_token_len, output.past_key_values
+        )
 
     def infer_stream(
         self,
         sample: datasets.VoiceSample,
-        max_tokens: Optional[int] = None,
+        max_new_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        num_beams: int = 1,
     ) -> base.InferenceGenerator:
         inputs = self._dataproc(sample)
         input_tokens = inputs["input_ids"].shape[1]
@@ -57,7 +68,14 @@ def infer_stream(
             self.tokenizer, skip_prompt=True, decode_kwargs=decode_kwargs
         )
 
-        thread_args = (inputs, max_tokens, temperature, streamer)
+        thread_args = (
+            inputs,
+            max_new_tokens,
+            temperature,
+            past_key_values,
+            num_beams,
+            streamer,
+        )
         thread = threading.Thread(target=self._generate, args=thread_args)
         thread.start()
         output_tokens = 0
@@ -108,8 +126,10 @@ def _dataproc(self, sample: datasets.VoiceSample):
     def _generate(
         self,
         inputs: torch.Tensor,
-        max_tokens: Optional[int] = None,
+        max_new_tokens: Optional[int] = None,
         temperature: Optional[float] = None,
+        past_key_values: Optional[Union[Tuple, transformers.cache_utils.Cache]] = None,
+        num_beams: int = 1,
         streamer: Optional[transformers.TextStreamer] = None,
     ):
         temperature = temperature or None
@@ -122,10 +142,13 @@ def _generate(
         return self.model.generate(
             **inputs,
             do_sample=do_sample,
-            max_new_tokens=max_tokens or MAX_TOKENS,
+            max_new_tokens=max_new_tokens or MAX_NEW_TOKENS,
             temperature=temperature,
             repetition_penalty=REPETITION_PENALTY,
             pad_token_id=self.tokenizer.eos_token_id,
             eos_token_id=terminators,
             streamer=streamer,
+            past_key_values=past_key_values,
+            num_beams=num_beams,
+            return_dict_in_generate=True,
         )

ultravox/model/ultravox_model.py

@@ -247,10 +247,16 @@ def prepare_inputs_for_generation(
             **kwargs,
         )
 
-        if is_cache_empty(past_key_values) and audio_values is not None:
-            # We only want to use audio features in the 1st generation step
+        prefill_start_idx = kwargs["cache_position"][0]
+        if (
+            audio_values is not None
+            and audio_token_start_idx is not None
+            and prefill_start_idx <= torch.max(audio_token_start_idx)
+        ):
             model_input["audio_values"] = audio_values
-            model_input["audio_token_start_idx"] = audio_token_start_idx
+            model_input["audio_token_start_idx"] = (
+                audio_token_start_idx - prefill_start_idx
+            )
             model_input["audio_token_len"] = audio_token_len
 
         return model_input