# Copyright (c) 2023 Graphcore Ltd. All rights reserved. from typing import Union, List, Optional, Any import abc import logging from functools import partial import torch import datasets import popart import popxl from popxl.utils import to_numpy from inference import inference from modelling.embedding import T5EmbeddingsTP from modelling.hf_mapping import hf_mapping_lm_tp from utils.utils import SimpleTimer from popxl_addons import timer from popxl_addons.array_munging import repeat, tensor_parallel_input from utils.inference import batch_inference from config import T5Config from datasets.utils.logging import disable_progress_bar from transformers.models.t5 import T5ForConditionalGeneration from transformers import AutoTokenizer class BasicPipeline: @abc.abstractmethod def __init__(self, *args, **kwargs) -> None: pass @abc.abstractmethod def __call__(self, prompt: Union[str, List[str]]) -> Any: pass def unwrap(dl): for example in dl: yield ( torch.tensor(example["input_ids"], dtype=torch.long), torch.tensor(example["attention_mask"], dtype=torch.float16), ) def encode_for_inference(tokenizer: AutoTokenizer): def func(dataset: datasets.Dataset): tokenized = tokenizer(dataset["text"], padding="max_length", return_tensors="np") return tokenized return func class T5Pipeline(BasicPipeline): def __init__( self, config: T5Config, model_size: str, hf_checkpoint: Union[str, T5ForConditionalGeneration], *args, sequence_length: Optional[int] = None, micro_batch_size: Optional[int] = None, output_length: Optional[int] = None, print_live: bool = True, tokenizer: Optional[AutoTokenizer] = None, **kwargs, ) -> None: """ A module which provides an interface to this application resembling the Hugging Face Transformers Pipeline. Args: - config: the inference configuration. - model_size: which model size to use, it can be either "xl" or "xxl". - pretrained: the Hugging Face pre-trained model, it can be the model's name, or the model object. - sequence_length: the maximum sequence length for the model. If provided, it will override the value in the configuration. - micro_batch_size: how many sequences the model will process in parallel. If provided, it will override the value in the configuration. - output_length: maximum number of tokens to generate. If provided, it will override the value in the configuration. - print_live: whether to print the new tokens as they're generated by the model. - tokenizer: the tokenizer. If not provided, one will be created based on the model's name. Note that `output_length` and `print_live` can also be overridden later on when calling this object. """ super().__init__(*args, **kwargs) if sequence_length is not None: config.model.sequence_length = sequence_length if micro_batch_size is not None: config.execution.micro_batch_size = micro_batch_size if output_length is not None: config.inference.output_length = output_length # Use float32 in order to prevent overflows config.model.dtype = popxl.float32 if model_size == "xl": # Reduce tensor parallel, so that it can be run on a POD4 config.execution.tensor_parallel = 4 logging.info(f"Creating session") session: popxl.Session = inference(config) if isinstance(hf_checkpoint, str): logging.info(f"Downloading '{hf_checkpoint}' pretrained weights") hf_model = T5ForConditionalGeneration.from_pretrained(hf_checkpoint) if tokenizer is None: logging.info(f"Downloading '{hf_checkpoint}' tokenizer") tokenizer = AutoTokenizer.from_pretrained(hf_checkpoint) else: hf_model = hf_checkpoint if tokenizer is None: raise ValueError( "Tokenizer needs to be passed to the Pipeline if a custom checkpoint is being provided." "Use: AutoTokenizer.from_pretrained(model-name) to create the tokenizer." ) with timer("Preparing HF pretrained model for loading to IPU"): weights = hf_mapping_lm_tp(config, session, hf_model) session.write_variables_data(weights) self.tokenizer = tokenizer self.config = config self.session = session self.output_length = config.inference.output_length self.print_live: bool = print_live self.tp = self.config.execution.tensor_parallel self.rf = self.config.execution.tensor_parallel * self.config.execution.data_parallel self.first_call = True def next_token(self, enc_batch, enc_masks, dec_batch, dec_masks, batch_lens): data_map = {} words = to_numpy(enc_batch, self.session.inputs.words.dtype).reshape(-1, *self.session.inputs.words.shape) attention_mask = to_numpy(enc_masks, self.session.inputs.attention_mask.dtype).reshape( -1, *self.session.inputs.attention_mask.shape ) decoder_words = to_numpy(dec_batch, self.session.inputs.decoder_words.dtype).reshape( -1, *self.session.inputs.decoder_words.shape ) decoder_attention_mask = to_numpy(dec_masks, self.session.inputs.decoder_attention_mask.dtype).reshape( -1, *self.session.inputs.decoder_attention_mask.shape ) data_map[self.session.inputs.words] = tensor_parallel_input( words, self.tp, self.rf, partial(T5EmbeddingsTP.offset_input, config=self.config) ) data_map[self.session.inputs.attention_mask] = tensor_parallel_input(attention_mask, self.tp, self.rf) data_map[self.session.inputs.decoder_words] = tensor_parallel_input( decoder_words, self.tp, self.rf, partial(T5EmbeddingsTP.offset_input, config=self.config) ) data_map[self.session.inputs.decoder_attention_mask] = tensor_parallel_input( decoder_attention_mask, self.tp, self.rf ) data_map[self.session.inputs.last_token_indices] = repeat(batch_lens - 1, self.tp, axis=0) # identical for all tp, take first with SimpleTimer() as simple_timer: next_token_id = self.session.run(data_map)[self.session.outputs.next_token][0] self.token_generation_time += simple_timer.elapsed self.token_generation_count += next_token_id.shape[0] if self.print_live: print(self.tokenizer.decode(next_token_id[0]), end=" ", flush=True) return torch.LongTensor(next_token_id) def __call__( self, prompt: Union[str, List[str], datasets.Dataset], *args, output_length: Optional[int] = None, print_live: Optional[bool] = None, ) -> List[str]: """ Perform inference with the given input prompt(s). Args: - prompt: the input(s) to pass to the model for inference. It can be a single string (single input), a list of strings (multiple inputs), or a dataset object, with the input prompts in `dataset["text"]`. - output_length: maximum number of tokens to generate. If provided, it will override the previous value. - print_live: whether to print the new tokens as they're generated by the model. If provided, it will override the previous value. Returns: - A list containing strings corresponding to the generated outputs, de-tokenised. """ super().__call__(prompt) if print_live is not None: self.print_live = print_live if output_length is not None: self.output_length = output_length micro_batch_size = self.config.execution.micro_batch_size # Preprocess the data including batching it if isinstance(prompt, str): prompt = [prompt] if isinstance(prompt, datasets.Dataset): data = prompt else: data = datasets.Dataset.from_dict({"text": prompt}) # No need to show the progress bar for this step disable_progress_bar() data = data.map( encode_for_inference(self.tokenizer), batched=True, remove_columns=data.column_names, load_from_cache_file=False, desc="Tokenizing prompt", ) if self.first_call: logging.info("Attaching and uploading weights to IPUs") self.first_call = False self.session.__enter__() if self.print_live: print(f"Prompt: '{prompt[0] if isinstance(prompt, list) else prompt[0]['text']}'") self.token_generation_time = 0 self.token_generation_count = 0 answers = batch_inference( unwrap(data), self.next_token, self.config.model.sequence_length, eos_token_id=self.tokenizer.eos_token_id, pad_token_id=self.tokenizer.pad_token_id, output_length=self.output_length, micro_batch_size=micro_batch_size, ) self.token_generation_time /= self.token_generation_count text_outputs = [self.tokenizer.decode(a, skip_special_tokens=True) for a in answers] return text_outputs def detach(self): was_attached_or_device = self.session._was_attached_stack.pop() self.session._device.detach() self.session._pb_session.setEngineIsLoaded(False) # If a DeviceInfo was stored in the stack then restore it. if isinstance(was_attached_or_device, popart.DeviceInfo): self.session._set_device(was_attached_or_device) self.first_call = True