nxphi47/MultiPurpose-Chatbot-DEMO

1 2import os3import numpy as np4import argparse5import torch6import gradio as gr7from typing import Any, Iterator8from typing import Iterator, List, Optional, Tuple9import filelock10import glob11import json12import time13from gradio.routes import Request14from gradio.utils import SyncToAsyncIterator, async_iteration15from gradio.helpers import special_args16import anyio17from typing import AsyncGenerator, Callable, Literal, Union, cast18 19from gradio_client.documentation import document, set_documentation_group20 21from typing import List, Optional, Union, Dict, Tuple22from tqdm.auto import tqdm23from huggingface_hub import snapshot_download24import types25 26from gradio.components import Button27from gradio.events import Dependency, EventListenerMethod28 29from .base_engine import BaseEngine30 31# ! Remember to use static cache32 33from transformers import (34    GenerationConfig,35    GenerationMixin,36    LogitsProcessorList,37    StoppingCriteriaList,38    DisjunctiveConstraint,39    BeamSearchScorer,40    PhrasalConstraint,41    ConstrainedBeamSearchScorer,42    PreTrainedModel,43)44import numpy as np45import random46import warnings47import inspect48from transformers.generation.utils import GenerateOutput, SampleOutput, logger49import torch50from typing import Callable, List, Optional, Union51from torch import nn52import torch.distributed as dist53import copy54 55from ..configs import (56    MODEL_PATH,57    DTYPE,58    DEVICE,59)60 61 62def setup_seed(seed):63    if seed == -1:64        return65    torch.manual_seed(seed)66    if torch.cuda.is_available():67        torch.cuda.manual_seed_all(seed)68    np.random.seed(seed)69    random.seed(seed)70    torch.backends.cudnn.deterministic = True71 72 73class NewGenerationMixin(GenerationMixin):74    """75    Allow generator sampling76 77    """78 79    # ! Copy from transformers.generation.utils -> GenerationMixin80    # Change sample function to sample_stream81    @torch.no_grad()82    def sample_stream(83        self,84        input_ids: torch.LongTensor,85        logits_processor: Optional[LogitsProcessorList] = None,86        stopping_criteria: Optional[StoppingCriteriaList] = None,87        logits_warper: Optional[LogitsProcessorList] = None,88        max_length: Optional[int] = None,89        pad_token_id: Optional[int] = None,90        eos_token_id: Optional[Union[int, List[int]]] = None,91        output_attentions: Optional[bool] = None,92        output_hidden_states: Optional[bool] = None,93        output_scores: Optional[bool] = None,94        output_logits: Optional[bool] = None,95        return_dict_in_generate: Optional[bool] = None,96        synced_gpus: bool = False,97        streamer: Optional["BaseStreamer"] = None,98        **model_kwargs,99    ):100        r"""101        Generates sequences of token ids for models with a language modeling head using **multinomial sampling** and102        can be used for text-decoder, text-to-text, speech-to-text, and vision-to-text models.103 104        <Tip warning={true}>105 106        In most cases, you do not need to call [`~generation.GenerationMixin.sample`] directly. Use generate() instead.107        For an overview of generation strategies and code examples, check the [following108        guide](../generation_strategies).109 110        </Tip>111 112        Parameters:113            input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):114                The sequence used as a prompt for the generation.115            logits_processor (`LogitsProcessorList`, *optional*):116                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsProcessor`]117                used to modify the prediction scores of the language modeling head applied at each generation step.118            stopping_criteria (`StoppingCriteriaList`, *optional*):119                An instance of [`StoppingCriteriaList`]. List of instances of class derived from [`StoppingCriteria`]120                used to tell if the generation loop should stop.121            logits_warper (`LogitsProcessorList`, *optional*):122                An instance of [`LogitsProcessorList`]. List of instances of class derived from [`LogitsWarper`] used123                to warp the prediction score distribution of the language modeling head applied before multinomial124                sampling at each generation step.125            max_length (`int`, *optional*, defaults to 20):126                **DEPRECATED**. Use `logits_processor` or `stopping_criteria` directly to cap the number of generated127                tokens. The maximum length of the sequence to be generated.128            pad_token_id (`int`, *optional*):129                The id of the *padding* token.130            eos_token_id (`Union[int, List[int]]`, *optional*):131                The id of the *end-of-sequence* token. Optionally, use a list to set multiple *end-of-sequence* tokens.132            output_attentions (`bool`, *optional*, defaults to `False`):133                Whether or not to return the attentions tensors of all attention layers. See `attentions` under134                returned tensors for more details.135            output_hidden_states (`bool`, *optional*, defaults to `False`):136                Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors137                for more details.138            output_scores (`bool`, *optional*, defaults to `False`):139                Whether or not to return the prediction scores. See `scores` under returned tensors for more details.140            output_logits (`bool`, *optional*, defaults to `False`):141                Whether or not to return the raw prediction logit scores. See `logits` under returned tensors for142                more details.143            return_dict_in_generate (`bool`, *optional*, defaults to `False`):144                Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.145            synced_gpus (`bool`, *optional*, defaults to `False`):146                Whether to continue running the while loop until max_length (needed for ZeRO stage 3)147            streamer (`BaseStreamer`, *optional*):148                Streamer object that will be used to stream the generated sequences. Generated tokens are passed149                through `streamer.put(token_ids)` and the streamer is responsible for any further processing.150            model_kwargs:151                Additional model specific kwargs will be forwarded to the `forward` function of the model. If model is152                an encoder-decoder model the kwargs should include `encoder_outputs`.153 154        Return:155            [`~generation.GenerateDecoderOnlyOutput`], [`~generation.GenerateEncoderDecoderOutput`] or `torch.LongTensor`:156            A `torch.LongTensor` containing the generated tokens (default behaviour) or a157            [`~generation.GenerateDecoderOnlyOutput`] if `model.config.is_encoder_decoder=False` and158            `return_dict_in_generate=True` or a [`~generation.GenerateEncoderDecoderOutput`] if159            `model.config.is_encoder_decoder=True`.160 161        Examples:162 163        ```python164        >>> from transformers import (165        ...     AutoTokenizer,166        ...     AutoModelForCausalLM,167        ...     LogitsProcessorList,168        ...     MinLengthLogitsProcessor,169        ...     TopKLogitsWarper,170        ...     TemperatureLogitsWarper,171        ...     StoppingCriteriaList,172        ...     MaxLengthCriteria,173        ... )174        >>> import torch175 176        >>> tokenizer = AutoTokenizer.from_pretrained("openai-community/gpt2")177        >>> model = AutoModelForCausalLM.from_pretrained("openai-community/gpt2")178 179        >>> # set pad_token_id to eos_token_id because GPT2 does not have a EOS token180        >>> model.config.pad_token_id = model.config.eos_token_id181        >>> model.generation_config.pad_token_id = model.config.eos_token_id182 183        >>> input_prompt = "Today is a beautiful day, and"184        >>> input_ids = tokenizer(input_prompt, return_tensors="pt").input_ids185 186        >>> # instantiate logits processors187        >>> logits_processor = LogitsProcessorList(188        ...     [189        ...         MinLengthLogitsProcessor(15, eos_token_id=model.generation_config.eos_token_id),190        ...     ]191        ... )192        >>> # instantiate logits processors193        >>> logits_warper = LogitsProcessorList(194        ...     [195        ...         TopKLogitsWarper(50),196        ...         TemperatureLogitsWarper(0.7),197        ...     ]198        ... )199 200        >>> stopping_criteria = StoppingCriteriaList([MaxLengthCriteria(max_length=20)])201 202        >>> torch.manual_seed(0)  # doctest: +IGNORE_RESULT203        >>> outputs = model.sample(204        ...     input_ids,205        ...     logits_processor=logits_processor,206        ...     logits_warper=logits_warper,207        ...     stopping_criteria=stopping_criteria,208        ... )209 210        >>> tokenizer.batch_decode(outputs, skip_special_tokens=True)211        ['Today is a beautiful day, and we must do everything possible to make it a day of celebration.']212        ```"""213        # init values214        from transformers.generation.utils import (215            validate_stopping_criteria, GenerateEncoderDecoderOutput, GenerateDecoderOnlyOutput216        )217        logits_processor = logits_processor if logits_processor is not None else LogitsProcessorList()218        stopping_criteria = stopping_criteria if stopping_criteria is not None else StoppingCriteriaList()219        if max_length is not None:220            warnings.warn(221                "`max_length` is deprecated in this function, use"222                " `stopping_criteria=StoppingCriteriaList([MaxLengthCriteria(max_length=max_length)])` instead.",223                UserWarning,224            )225            stopping_criteria = validate_stopping_criteria(stopping_criteria, max_length)226        logits_warper = logits_warper if logits_warper is not None else LogitsProcessorList()227        pad_token_id = pad_token_id if pad_token_id is not None else self.generation_config.pad_token_id228        eos_token_id = eos_token_id if eos_token_id is not None else self.generation_config.eos_token_id229        if isinstance(eos_token_id, int):230            eos_token_id = [eos_token_id]231        eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None232        output_scores = output_scores if output_scores is not None else self.generation_config.output_scores233        output_logits = output_logits if output_logits is not None else self.generation_config.output_logits234        output_attentions = (235            output_attentions if output_attentions is not None else self.generation_config.output_attentions236        )237        output_hidden_states = (238            output_hidden_states if output_hidden_states is not None else self.generation_config.output_hidden_states239        )240        return_dict_in_generate = (241            return_dict_in_generate242            if return_dict_in_generate is not None243            else self.generation_config.return_dict_in_generate244        )245 246        # init attention / hidden states / scores tuples247        scores = () if (return_dict_in_generate and output_scores) else None248        raw_logits = () if (return_dict_in_generate and output_logits) else None249        decoder_attentions = () if (return_dict_in_generate and output_attentions) else None250        cross_attentions = () if (return_dict_in_generate and output_attentions) else None251        decoder_hidden_states = () if (return_dict_in_generate and output_hidden_states) else None252 253        # if model is an encoder-decoder, retrieve encoder attention weights and hidden states254        if return_dict_in_generate and self.config.is_encoder_decoder:255            encoder_attentions = model_kwargs["encoder_outputs"].get("attentions") if output_attentions else None256            encoder_hidden_states = (257                model_kwargs["encoder_outputs"].get("hidden_states") if output_hidden_states else None258            )259        # keep track of which sequences are already finished260        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)261 262        this_peer_finished = False  # used by synced_gpus only263        # auto-regressive generation264        while True:265            if synced_gpus:266                # Under synced_gpus the `forward` call must continue until all gpus complete their sequence.267                # The following logic allows an early break if all peers finished generating their sequence268                this_peer_finished_flag = torch.tensor(0.0 if this_peer_finished else 1.0).to(input_ids.device)269                # send 0.0 if we finished, 1.0 otherwise270                dist.all_reduce(this_peer_finished_flag, op=dist.ReduceOp.SUM)271                # did all peers finish? the reduced sum will be 0.0 then272                if this_peer_finished_flag.item() == 0.0:273                    break274 275            # prepare model inputs276            model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)277 278            # forward pass to get next token279            outputs = self(280                **model_inputs,281                return_dict=True,282                output_attentions=output_attentions,283                output_hidden_states=output_hidden_states,284            )285 286            if synced_gpus and this_peer_finished:287                continue  # don't waste resources running the code we don't need288 289            next_token_logits = outputs.logits[:, -1, :]290 291            # pre-process distribution292            next_token_scores = logits_processor(input_ids, next_token_logits)293            next_token_scores = logits_warper(input_ids, next_token_scores)294 295            # Store scores, attentions and hidden_states when required296            if return_dict_in_generate:297                if output_scores:298                    scores += (next_token_scores,)299                if output_logits:300                    raw_logits += (next_token_logits,)301                if output_attentions:302                    decoder_attentions += (303                        (outputs.decoder_attentions,) if self.config.is_encoder_decoder else (outputs.attentions,)304                    )305                    if self.config.is_encoder_decoder:306                        cross_attentions += (outputs.cross_attentions,)307 308                if output_hidden_states:309                    decoder_hidden_states += (310                        (outputs.decoder_hidden_states,)311                        if self.config.is_encoder_decoder312                        else (outputs.hidden_states,)313                    )314 315            # sample316            probs = nn.functional.softmax(next_token_scores, dim=-1)317            next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1)318 319            # finished sentences should have their next token be a padding token320            if eos_token_id is not None:321                if pad_token_id is None:322                    raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.")323                next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences)324 325            yield next_tokens.cpu()326 327            # update generated ids, model inputs, and length for next step328            input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1)329            if streamer is not None:330                streamer.put(next_tokens.cpu())331            332            next_model_inputs = {}333            if "cache_position" in model_inputs:334                next_model_inputs['cache_position'] = model_inputs['cache_position']335            336            try:337                model_kwargs = self._update_model_kwargs_for_generation(338                    outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder, 339                    # model_inputs=model_inputs340                    model_inputs=next_model_inputs,341                )342            except Exception as e:343                # Older version dont have model_inputs344                model_kwargs = self._update_model_kwargs_for_generation(345                    outputs, model_kwargs, is_encoder_decoder=self.config.is_encoder_decoder, 346                )347 348 349            # if eos_token was found in one sentence, set sentence to finished350            if eos_token_id_tensor is not None:351                unfinished_sequences = unfinished_sequences.mul(352                    next_tokens.tile(eos_token_id_tensor.shape[0], 1).ne(eos_token_id_tensor.unsqueeze(1)).prod(dim=0)353                )354 355                # stop when each sentence is finished356                if unfinished_sequences.max() == 0:357                    this_peer_finished = True358 359            # stop if we exceed the maximum length360            if stopping_criteria(input_ids, scores):361                this_peer_finished = True362 363            if this_peer_finished and not synced_gpus:364                break365 366        if streamer is not None:367            streamer.end()368 369        # if return_dict_in_generate:370        #     if self.config.is_encoder_decoder:371        #         return GenerateEncoderDecoderOutput(372        #             sequences=input_ids,373        #             scores=scores,374        #             logits=raw_logits,375        #             encoder_attentions=encoder_attentions,376        #             encoder_hidden_states=encoder_hidden_states,377        #             decoder_attentions=decoder_attentions,378        #             cross_attentions=cross_attentions,379        #             decoder_hidden_states=decoder_hidden_states,380        #             past_key_values=model_kwargs.get("past_key_values"),381        #         )382        #     else:383        #         return GenerateDecoderOnlyOutput(384        #             sequences=input_ids,385        #             scores=scores,386        #             logits=raw_logits,387        #             attentions=decoder_attentions,388        #             hidden_states=decoder_hidden_states,389        #             past_key_values=model_kwargs.get("past_key_values"),390        #         )391        # else:392        #     return input_ids393 394 395 396class TransformersEngine(BaseEngine):397    @property398    def max_position_embeddings(self) -> int:399        return self._model.config.max_position_embeddings400 401    @property402    def tokenizer(self):403        return self._tokenizer404 405    def load_model(self):406        from transformers import AutoTokenizer, AutoModelForCausalLM407        import sys408        # caution: path[0] is reserved for script path (or '' in REPL)409        # sys.path.append(CODE_PATH)410        self.model_path = model_path = MODEL_PATH411        self.torch_dtype = torch.bfloat16 if DTYPE == 'bfloat16' else torch.float16412        self.device_map = DEVICE413        print(f'Loading model from {model_path} on {self.device_map} with {self.torch_dtype}')414 415        self._tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)416        assert self._tokenizer.chat_template is not None and self._tokenizer.chat_template != "", f"{self._tokenizer.chat_template=} not found!"417        self._model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=self.torch_dtype, device_map=self.device_map, trust_remote_code=True).eval()418        self._model.sample_old = self._model.sample419        self._model._sample = types.MethodType(NewGenerationMixin.sample_stream, self._model)420        print(self._model)421        print(f"{self.max_position_embeddings=}")422    423    def generate_yield_string(self, prompt, temperature, max_tokens, stop_strings: Optional[Tuple[str]] = None, **kwargs):424        425        # ! MUST PUT INSIDE torch.no_grad() otherwise it will overflow OOM426        with torch.no_grad():427            inputs = self.tokenizer(prompt, return_tensors='pt')428            num_tokens = inputs.input_ids.size(1)429 430            inputs = {k: v.to(self.device_map) for k, v in inputs.items() if v is not None}431            generator = self._model.generate(432                **inputs, 433                do_sample=True, 434                temperature=temperature, 435                max_new_tokens=max_tokens, 436                pad_token_id=self.processor.tokenizer.pad_token_id,437            )438 439            out_tokens = []440            response = None441            for token in generator:442                out_tokens.append(token.item())443                response = self.processor.tokenizer.decode(out_tokens)444                num_tokens += 1445                # print(f"{num_tokens=}", end='\r')446                # sys.stdout.flush()447                yield response, num_tokens448            449            if response is not None:450                full_text = prompt + response451                num_tokens = len(self.tokenizer.encode(full_text))452                yield response, num_tokens