|
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) |
|
|
