| 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 |
|
|
| logger = logging.getLogger(__name__) |
|
|
| console = Console() |
|
|
|
|
| WHISPER_LANGUAGE_TO_LLM_LANGUAGE = { |
| "en": "english", |
| "fr": "french", |
| "es": "spanish", |
| "zh": "chinese", |
| "ja": "japanese", |
| "ko": "korean", |
| } |
|
|
| 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 = "Repeat the word 'home'." |
| 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): |
| console.print("infering language model...") |
| console.print(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}) |
| 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], language_code) |
| printable_text = new_text |
|
|
| self.chat.append({"role": "assistant", "content": generated_text}) |
|
|
| |
| yield (printable_text, language_code) |
|
|
