| | import logging |
| | from LLM.chat import Chat |
| | from baseHandler import BaseHandler |
| | from mlx_lm import load, stream_generate, generate |
| | from rich.console import Console |
| | import torch |
| |
|
| | logger = logging.getLogger(__name__) |
| |
|
| | console = Console() |
| |
|
| | WHISPER_LANGUAGE_TO_LLM_LANGUAGE = { |
| | "en": "english", |
| | "fr": "french", |
| | "es": "spanish", |
| | "zh": "chinese", |
| | "ja": "japanese", |
| | "ko": "korean", |
| | } |
| |
|
| | class MLXLanguageModelHandler(BaseHandler): |
| | """ |
| | Handles the language model part. |
| | """ |
| |
|
| | def setup( |
| | self, |
| | model_name="microsoft/Phi-3-mini-4k-instruct", |
| | device="mps", |
| | 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.model_name = model_name |
| | self.model, self.tokenizer = load(self.model_name) |
| | self.gen_kwargs = 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 = "Repeat the word 'home'." |
| | dummy_chat = [{"role": self.user_role, "content": dummy_input_text}] |
| |
|
| | n_steps = 2 |
| |
|
| | for _ in range(n_steps): |
| | prompt = self.tokenizer.apply_chat_template(dummy_chat, tokenize=False) |
| | generate( |
| | self.model, |
| | self.tokenizer, |
| | prompt=prompt, |
| | max_tokens=self.gen_kwargs["max_new_tokens"], |
| | verbose=False, |
| | ) |
| |
|
| | def process(self, prompt): |
| | logger.debug("infering language model...") |
| | language_code = None |
| |
|
| | if isinstance(prompt, tuple): |
| | prompt, language_code = prompt |
| | prompt = f"Please reply to my message in {WHISPER_LANGUAGE_TO_LLM_LANGUAGE[language_code]}. " + prompt |
| |
|
| | self.chat.append({"role": self.user_role, "content": prompt}) |
| |
|
| | |
| | if "gemma" in self.model_name.lower(): |
| | chat_messages = [ |
| | msg for msg in self.chat.to_list() if msg["role"] != "system" |
| | ] |
| | else: |
| | chat_messages = self.chat.to_list() |
| |
|
| | prompt = self.tokenizer.apply_chat_template( |
| | chat_messages, tokenize=False, add_generation_prompt=True |
| | ) |
| | output = "" |
| | curr_output = "" |
| | for t in stream_generate( |
| | self.model, |
| | self.tokenizer, |
| | prompt, |
| | max_tokens=self.gen_kwargs["max_new_tokens"], |
| | ): |
| | output += t |
| | curr_output += t |
| | if curr_output.endswith((".", "?", "!", "<|end|>")): |
| | yield (curr_output.replace("<|end|>", ""), language_code) |
| | curr_output = "" |
| | generated_text = output.replace("<|end|>", "") |
| | torch.mps.empty_cache() |
| |
|
| | self.chat.append({"role": "assistant", "content": generated_text}) |
