Spaces:
Running
Running
| import functools | |
| from src.enums import t5_type | |
| def get_loaders(model_name, reward_type, llama_type=None, load_gptq='', load_exllama=False, config=None, | |
| rope_scaling=None, max_seq_len=None, model_name_exllama_if_no_config=''): | |
| # NOTE: Some models need specific new prompt_type | |
| # E.g. t5_xxl_true_nli_mixture has input format: "premise: PREMISE_TEXT hypothesis: HYPOTHESIS_TEXT".) | |
| if load_exllama: | |
| from src.llm_exllama import H2OExLlamaTokenizer, H2OExLlamaGenerator | |
| from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig | |
| import os, glob | |
| if config: | |
| # then use HF path | |
| from transformers import TRANSFORMERS_CACHE | |
| model_directory = os.path.join(TRANSFORMERS_CACHE, 'models--' + config.name_or_path.replace('/', '--'), | |
| 'snapshots', config._commit_hash) | |
| else: | |
| # then use path in env file | |
| # Directory containing model, tokenizer, generator | |
| model_directory = model_name_exllama_if_no_config | |
| # download model | |
| revision = config._commit_hash | |
| from huggingface_hub import snapshot_download | |
| snapshot_download(repo_id=model_name, revision=revision) | |
| # Locate files we need within that directory | |
| tokenizer_path = os.path.join(model_directory, "tokenizer.model") | |
| assert os.path.isfile(tokenizer_path), "Missing %s" % tokenizer_path | |
| model_config_path = os.path.join(model_directory, "config.json") | |
| assert os.path.isfile(model_config_path), "Missing %s" % model_config_path | |
| st_pattern = os.path.join(model_directory, "*.safetensors") | |
| model_path = glob.glob(st_pattern)[0] | |
| assert os.path.isfile(model_path), "Missing %s" % model_path | |
| # Create config, model, tokenizer and generator | |
| exconfig = ExLlamaConfig(model_config_path) # create config from config.json | |
| rope_scaling = rope_scaling or {} | |
| exconfig.alpha_value = rope_scaling.get('alpha_value', 1) # rope | |
| exconfig.compress_pos_emb = rope_scaling.get('compress_pos_emb', 1) # related rope | |
| # update max_seq_len | |
| assert hasattr(config, 'max_position_embeddings') or hasattr(config, | |
| 'max_sequence_length'), "Improve code if no such argument" | |
| if hasattr(config, 'max_position_embeddings'): | |
| exconfig.max_seq_len = int(config.max_position_embeddings * exconfig.alpha_value) | |
| else: | |
| exconfig.max_seq_len = int(config.max_sequence_length * exconfig.alpha_value) | |
| if 'Llama-2'.lower() in model_name.lower(): | |
| # override bad defaults | |
| exconfig.max_seq_len = int(4096 * exconfig.alpha_value) | |
| if max_seq_len is not None: | |
| exconfig.max_seq_len = max_seq_len | |
| exconfig.model_path = model_path # supply path to model weights file | |
| model = ExLlama(exconfig) # create ExLlama instance and load the weights | |
| tokenizer = H2OExLlamaTokenizer(tokenizer_path) # create tokenizer from tokenizer model file | |
| tokenizer.model_max_length = exconfig.max_seq_len | |
| cache = ExLlamaCache(model) # create cache for inference | |
| generator = H2OExLlamaGenerator(model, tokenizer, cache) # create generator | |
| return generator, tokenizer, False | |
| if load_gptq: | |
| from transformers import AutoTokenizer | |
| from auto_gptq import AutoGPTQForCausalLM | |
| use_triton = False | |
| model_loader = functools.partial(AutoGPTQForCausalLM.from_quantized, | |
| quantize_config=None, use_triton=use_triton, | |
| ) | |
| return model_loader, AutoTokenizer, False | |
| if llama_type is None: | |
| llama_type = "llama" in model_name.lower() | |
| if llama_type: | |
| from transformers import LlamaForCausalLM, LlamaTokenizer | |
| return LlamaForCausalLM.from_pretrained, LlamaTokenizer, False | |
| elif 'distilgpt2' in model_name.lower(): | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| return AutoModelForCausalLM.from_pretrained, AutoTokenizer, False | |
| elif 'gpt2' in model_name.lower(): | |
| from transformers import GPT2LMHeadModel, GPT2Tokenizer | |
| return GPT2LMHeadModel.from_pretrained, GPT2Tokenizer, False | |
| elif 'mbart-' in model_name.lower(): | |
| from transformers import MBartForConditionalGeneration, MBart50TokenizerFast | |
| return MBartForConditionalGeneration.from_pretrained, MBart50TokenizerFast, True | |
| elif t5_type(model_name): | |
| from transformers import AutoTokenizer, T5ForConditionalGeneration | |
| return T5ForConditionalGeneration.from_pretrained, AutoTokenizer, True | |
| elif 'bigbird' in model_name: | |
| from transformers import BigBirdPegasusForConditionalGeneration, AutoTokenizer | |
| return BigBirdPegasusForConditionalGeneration.from_pretrained, AutoTokenizer, True | |
| elif 'bart-large-cnn-samsum' in model_name or 'flan-t5-base-samsum' in model_name: | |
| from transformers import pipeline | |
| return pipeline, "summarization", False | |
| elif reward_type or 'OpenAssistant/reward-model'.lower() in model_name.lower(): | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| return AutoModelForSequenceClassification.from_pretrained, AutoTokenizer, False | |
| else: | |
| from transformers import AutoTokenizer, AutoModelForCausalLM | |
| model_loader = AutoModelForCausalLM | |
| tokenizer_loader = AutoTokenizer | |
| return model_loader.from_pretrained, tokenizer_loader, False | |
| def get_tokenizer(tokenizer_loader, tokenizer_base_model, local_files_only, resume_download, use_auth_token): | |
| tokenizer = tokenizer_loader.from_pretrained(tokenizer_base_model, | |
| local_files_only=local_files_only, | |
| resume_download=resume_download, | |
| use_auth_token=use_auth_token, | |
| padding_side='left') | |
| tokenizer.pad_token_id = 0 # different from the eos token | |
| # when generating, we will use the logits of right-most token to predict the next token | |
| # so the padding should be on the left, | |
| # e.g. see: https://huggingface.co/transformers/v4.11.3/model_doc/t5.html#inference | |
| tokenizer.padding_side = "left" # Allow batched inference | |
| return tokenizer | |