Upload 5 files

LLM/chat.py

@@ -0,0 +1,25 @@
+class Chat:
+    """
+    Handles the chat using to avoid OOM issues.
+    """
+
+    def __init__(self, size):
+        self.size = size
+        self.init_chat_message = None
+        # maxlen is necessary pair, since a each new step we add an prompt and assitant answer
+        self.buffer = []
+
+    def append(self, item):
+        self.buffer.append(item)
+        if len(self.buffer) == 2 * (self.size + 1):
+            self.buffer.pop(0)
+            self.buffer.pop(0)
+
+    def init_chat(self, init_chat_message):
+        self.init_chat_message = init_chat_message
+
+    def to_list(self):
+        if self.init_chat_message:
+            return [self.init_chat_message] + self.buffer
+        else:
+            return self.buffer

LLM/language_model.py

@@ -0,0 +1,134 @@
+from threading import Thread
+from transformers import (
+    AutoModelForCausalLM,
+    AutoTokenizer,
+    pipeline,
+    TextIteratorStreamer,
+)
+import torch
+
+from LLM.chat import Chat
+from baseHandler import BaseHandler
+from rich.console import Console
+import logging
+from nltk import sent_tokenize
+
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+
+class LanguageModelHandler(BaseHandler):
+    """
+    Handles the language model part.
+    """
+
+    def setup(
+        self,
+        model_name="microsoft/Phi-3-mini-4k-instruct",
+        device="cuda",
+        torch_dtype="float16",
+        gen_kwargs={},
+        user_role="user",
+        chat_size=1,
+        init_chat_role=None,
+        init_chat_prompt="You are a helpful AI assistant.",
+    ):
+        self.device = device
+        self.torch_dtype = getattr(torch, torch_dtype)
+
+        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
+        self.model = AutoModelForCausalLM.from_pretrained(
+            model_name, torch_dtype=torch_dtype, trust_remote_code=True
+        ).to(device)
+        self.pipe = pipeline(
+            "text-generation", model=self.model, tokenizer=self.tokenizer, device=device
+        )
+        self.streamer = TextIteratorStreamer(
+            self.tokenizer,
+            skip_prompt=True,
+            skip_special_tokens=True,
+        )
+        self.gen_kwargs = {
+            "streamer": self.streamer,
+            "return_full_text": False,
+            **gen_kwargs,
+        }
+
+        self.chat = Chat(chat_size)
+        if init_chat_role:
+            if not init_chat_prompt:
+                raise ValueError(
+                    "An initial promt needs to be specified when setting init_chat_role."
+                )
+            self.chat.init_chat({"role": init_chat_role, "content": init_chat_prompt})
+        self.user_role = user_role
+
+        self.warmup()
+
+    def warmup(self):
+        logger.info(f"Warming up {self.__class__.__name__}")
+
+        dummy_input_text = "Write me a poem about Machine Learning."
+        dummy_chat = [{"role": self.user_role, "content": dummy_input_text}]
+        warmup_gen_kwargs = {
+            "min_new_tokens": self.gen_kwargs["min_new_tokens"],
+            "max_new_tokens": self.gen_kwargs["max_new_tokens"],
+            **self.gen_kwargs,
+        }
+
+        n_steps = 2
+
+        if self.device == "cuda":
+            start_event = torch.cuda.Event(enable_timing=True)
+            end_event = torch.cuda.Event(enable_timing=True)
+            torch.cuda.synchronize()
+            start_event.record()
+
+        for _ in range(n_steps):
+            thread = Thread(
+                target=self.pipe, args=(dummy_chat,), kwargs=warmup_gen_kwargs
+            )
+            thread.start()
+            for _ in self.streamer:
+                pass
+
+        if self.device == "cuda":
+            end_event.record()
+            torch.cuda.synchronize()
+
+            logger.info(
+                f"{self.__class__.__name__}:  warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s"
+            )
+
+    def process(self, prompt):
+        logger.debug("infering language model...")
+
+        self.chat.append({"role": self.user_role, "content": prompt})
+        thread = Thread(
+            target=self.pipe, args=(self.chat.to_list(),), kwargs=self.gen_kwargs
+        )
+        thread.start()
+        if self.device == "mps":
+            generated_text = ""
+            for new_text in self.streamer:
+                generated_text += new_text
+            printable_text = generated_text
+            torch.mps.empty_cache()
+        else:
+            generated_text, printable_text = "", ""
+            for new_text in self.streamer:
+                generated_text += new_text
+                printable_text += new_text
+                sentences = sent_tokenize(printable_text)
+                if len(sentences) > 1:
+                    yield (sentences[0])
+                    printable_text = new_text
+
+        self.chat.append({"role": "assistant", "content": generated_text})
+
+        # don't forget last sentence
+        yield printable_text

STT/lightning_whisper_mlx_handler.py

@@ -0,0 +1,58 @@
+import logging
+from time import perf_counter
+from baseHandler import BaseHandler
+from lightning_whisper_mlx import LightningWhisperMLX
+import numpy as np
+from rich.console import Console
+import torch
+
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+
+class LightningWhisperSTTHandler(BaseHandler):
+    """
+    Handles the Speech To Text generation using a Whisper model.
+    """
+
+    def setup(
+        self,
+        model_name="distil-large-v3",
+        device="cuda",
+        torch_dtype="float16",
+        compile_mode=None,
+        gen_kwargs={},
+    ):
+        if len(model_name.split("/")) > 1:
+            model_name = model_name.split("/")[-1]
+        self.device = device
+        self.model = LightningWhisperMLX(model=model_name, batch_size=6, quant=None)
+        self.warmup()
+
+    def warmup(self):
+        logger.info(f"Warming up {self.__class__.__name__}")
+
+        # 2 warmup steps for no compile or compile mode with CUDA graphs capture
+        n_steps = 1
+        dummy_input = np.array([0] * 512)
+
+        for _ in range(n_steps):
+            _ = self.model.transcribe(dummy_input)["text"].strip()
+
+    def process(self, spoken_prompt):
+        logger.debug("infering whisper...")
+
+        global pipeline_start
+        pipeline_start = perf_counter()
+
+        pred_text = self.model.transcribe(spoken_prompt)["text"].strip()
+        torch.mps.empty_cache()
+
+        logger.debug("finished whisper inference")
+        console.print(f"[yellow]USER: {pred_text}")
+
+        yield pred_text

STT/whisper_stt_handler.py

@@ -0,0 +1,113 @@
+from time import perf_counter
+from transformers import (
+    AutoModelForSpeechSeq2Seq,
+    AutoProcessor,
+)
+import torch
+
+from baseHandler import BaseHandler
+from rich.console import Console
+import logging
+
+logging.basicConfig(
+    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+)
+logger = logging.getLogger(__name__)
+
+console = Console()
+
+
+class WhisperSTTHandler(BaseHandler):
+    """
+    Handles the Speech To Text generation using a Whisper model.
+    """
+
+    def setup(
+        self,
+        model_name="distil-whisper/distil-large-v3",
+        device="cuda",
+        torch_dtype="float16",
+        compile_mode=None,
+        gen_kwargs={},
+    ):
+        self.device = device
+        self.torch_dtype = getattr(torch, torch_dtype)
+        self.compile_mode = compile_mode
+        self.gen_kwargs = gen_kwargs
+
+        self.processor = AutoProcessor.from_pretrained(model_name)
+        self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
+            model_name,
+            torch_dtype=self.torch_dtype,
+        ).to(device)
+
+        # compile
+        if self.compile_mode:
+            self.model.generation_config.cache_implementation = "static"
+            self.model.forward = torch.compile(
+                self.model.forward, mode=self.compile_mode, fullgraph=True
+            )
+        self.warmup()
+
+    def prepare_model_inputs(self, spoken_prompt):
+        input_features = self.processor(
+            spoken_prompt, sampling_rate=16000, return_tensors="pt"
+        ).input_features
+        input_features = input_features.to(self.device, dtype=self.torch_dtype)
+
+        return input_features
+
+    def warmup(self):
+        logger.info(f"Warming up {self.__class__.__name__}")
+
+        # 2 warmup steps for no compile or compile mode with CUDA graphs capture
+        n_steps = 1 if self.compile_mode == "default" else 2
+        dummy_input = torch.randn(
+            (1, self.model.config.num_mel_bins, 3000),
+            dtype=self.torch_dtype,
+            device=self.device,
+        )
+        if self.compile_mode not in (None, "default"):
+            # generating more tokens than previously will trigger CUDA graphs capture
+            # one should warmup with a number of generated tokens above max tokens targeted for subsequent generation
+            warmup_gen_kwargs = {
+                "min_new_tokens": self.gen_kwargs["min_new_tokens"],
+                "max_new_tokens": self.gen_kwargs["max_new_tokens"],
+                **self.gen_kwargs,
+            }
+        else:
+            warmup_gen_kwargs = self.gen_kwargs
+
+        if self.device == "cuda":
+            start_event = torch.cuda.Event(enable_timing=True)
+            end_event = torch.cuda.Event(enable_timing=True)
+            torch.cuda.synchronize()
+            start_event.record()
+
+        for _ in range(n_steps):
+            _ = self.model.generate(dummy_input, **warmup_gen_kwargs)
+
+        if self.device == "cuda":
+            end_event.record()
+            torch.cuda.synchronize()
+
+            logger.info(
+                f"{self.__class__.__name__}:  warmed up! time: {start_event.elapsed_time(end_event) * 1e-3:.3f} s"
+            )
+
+    def process(self, spoken_prompt):
+        logger.debug("infering whisper...")
+
+        global pipeline_start
+        pipeline_start = perf_counter()
+
+        input_features = self.prepare_model_inputs(spoken_prompt)
+        pred_ids = self.model.generate(input_features, **self.gen_kwargs)
+        pred_text = self.processor.batch_decode(
+            pred_ids, skip_special_tokens=True, decode_with_timestamps=False
+        )[0]
+
+        logger.debug("finished whisper inference")
+        console.print(f"[yellow]USER: {pred_text}")
+
+        yield pred_text

baseHandler.py

@@ -0,0 +1,51 @@
+from time import perf_counter
+import logging
+
+logger = logging.getLogger(__name__)
+
+
+class BaseHandler:
+    """
+    Base class for pipeline parts. Each part of the pipeline has an input and an output queue.
+    The `setup` method along with `setup_args` and `setup_kwargs` can be used to address the specific requirements of the implemented pipeline part.
+    To stop a handler properly, set the stop_event and, to avoid queue deadlocks, place b"END" in the input queue.
+    Objects placed in the input queue will be processed by the `process` method, and the yielded results will be placed in the output queue.
+    The cleanup method handles stopping the handler, and b"END" is placed in the output queue.
+    """
+
+    def __init__(self, stop_event, queue_in, queue_out, setup_args=(), setup_kwargs={}):
+        self.stop_event = stop_event
+        self.queue_in = queue_in
+        self.queue_out = queue_out
+        self.setup(*setup_args, **setup_kwargs)
+        self._times = []
+
+    def setup(self):
+        pass
+
+    def process(self):
+        raise NotImplementedError
+
+    def run(self):
+        while not self.stop_event.is_set():
+            input = self.queue_in.get()
+            if isinstance(input, bytes) and input == b"END":
+                # sentinelle signal to avoid queue deadlock
+                logger.debug("Stopping thread")
+                break
+            start_time = perf_counter()
+            for output in self.process(input):
+                self._times.append(perf_counter() - start_time)
+                logger.debug(f"{self.__class__.__name__}: {self.last_time: .3f} s")
+                self.queue_out.put(output)
+                start_time = perf_counter()
+
+        self.cleanup()
+        self.queue_out.put(b"END")
+
+    @property
+    def last_time(self):
+        return self._times[-1]
+
+    def cleanup(self):
+        pass