import os import sys import threading from queue import Queue from typing import Any, Optional import torch from langchain import HuggingFaceTextGenInference from langchain.callbacks.base import BaseCallbackHandler from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler from langchain.chat_models import ChatOpenAI from langchain.llms import ( CTransformers, GPT4All, HuggingFacePipeline, LlamaCpp, OpenLLM, ) from langchain.schema import LLMResult from transformers import ( AutoConfig, AutoModelForCausalLM, AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig, StoppingCriteria, StoppingCriteriaList, T5Tokenizer, TextStreamer, pipeline, ) from app_modules.instruct_pipeline import InstructionTextGenerationPipeline from app_modules.utils import ensure_model_is_downloaded class TextIteratorStreamer(TextStreamer, StreamingStdOutCallbackHandler): def __init__( self, tokenizer: "AutoTokenizer", skip_prompt: bool = False, timeout: Optional[float] = None, for_huggingface: bool = False, **decode_kwargs, ): super().__init__(tokenizer, skip_prompt, **decode_kwargs) self.text_queue = Queue() self.stop_signal = None self.timeout = timeout self.total_tokens = 0 self.for_huggingface = for_huggingface def on_finalized_text(self, text: str, stream_end: bool = False): super().on_finalized_text(text, stream_end=stream_end) """Put the new text in the queue. If the stream is ending, also put a stop signal in the queue.""" self.text_queue.put(text, timeout=self.timeout) self.total_tokens = self.total_tokens + 1 if stream_end: print("\n") self.text_queue.put("\n", timeout=self.timeout) self.text_queue.put(self.stop_signal, timeout=self.timeout) def on_llm_new_token(self, token: str, **kwargs: Any) -> None: sys.stdout.write(token) sys.stdout.flush() self.text_queue.put(token, timeout=self.timeout) self.total_tokens = self.total_tokens + 1 def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None: print("\n") self.text_queue.put("\n", timeout=self.timeout) self.text_queue.put(self.stop_signal, timeout=self.timeout) def __iter__(self): return self def __next__(self): value = self.text_queue.get(timeout=self.timeout) if value == self.stop_signal: raise StopIteration() else: return value def reset(self, q: Queue = None): # print("resetting TextIteratorStreamer") self.text_queue = q if q is not None else Queue() def empty(self): return self.text_queue.empty() class LLMLoader: llm_model_type: str llm: any streamer: any max_tokens_limit: int lock: any def __init__(self, llm_model_type): self.llm_model_type = llm_model_type self.llm = None self.streamer = TextIteratorStreamer( "", for_huggingface=True, ) self.max_tokens_limit = 2048 self.search_kwargs = {"k": 4} self.lock = threading.Lock() def _init_hf_streamer(self, tokenizer): self.streamer = TextIteratorStreamer( tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True, for_huggingface=True, ) def init( self, custom_handler: Optional[BaseCallbackHandler] = None, n_threds: int = 4, hf_pipeline_device_type: str = None, ): print("initializing LLM: " + self.llm_model_type) if hf_pipeline_device_type is None: hf_pipeline_device_type = "cpu" using_cuda = hf_pipeline_device_type.startswith("cuda") torch_dtype = torch.float16 if using_cuda else torch.float32 if os.environ.get("USING_TORCH_BFLOAT16") == "true": torch_dtype = torch.bfloat16 load_quantized_model = os.environ.get("LOAD_QUANTIZED_MODEL") print(f" hf_pipeline_device_type: {hf_pipeline_device_type}") print(f" load_quantized_model: {load_quantized_model}") print(f" torch_dtype: {torch_dtype}") print(f" n_threds: {n_threds}") double_quant_config = BitsAndBytesConfig( load_in_4bit=load_quantized_model == "4bit", bnb_4bit_use_double_quant=load_quantized_model == "4bit", load_in_8bit=load_quantized_model == "8bit", bnb_8bit_use_double_quant=load_quantized_model == "8bit", ) callbacks = [] if self.streamer is not None: callbacks.append(self.streamer) if custom_handler is not None: callbacks.append(custom_handler) if self.llm is None: if self.llm_model_type == "openai": MODEL_NAME = os.environ.get("OPENAI_MODEL_NAME") or "gpt-3.5-turbo" print(f" using model: {MODEL_NAME}") self.llm = ChatOpenAI( model_name=MODEL_NAME, streaming=True, callbacks=callbacks, verbose=True, temperature=0, ) elif self.llm_model_type == "openllm": server_url = os.environ.get("OPENLLM_SERVER_URL") print(f" server url: {server_url}") self.llm = OpenLLM( server_url=server_url, # callbacks=callbacks, verbose=True, ) elif self.llm_model_type.startswith("gpt4all"): MODEL_PATH = ensure_model_is_downloaded(self.llm_model_type) self.llm = GPT4All( model=MODEL_PATH, max_tokens=2048, n_threads=n_threds, backend="gptj" if self.llm_model_type == "gpt4all-j" else "llama", callbacks=callbacks, verbose=True, use_mlock=True, ) elif self.llm_model_type == "llamacpp": MODEL_PATH = ensure_model_is_downloaded(self.llm_model_type) self.llm = LlamaCpp( model_path=MODEL_PATH, n_ctx=8192, n_threads=n_threds, seed=0, temperature=0, max_tokens=2048, callbacks=callbacks, verbose=True, use_mlock=True, ) elif self.llm_model_type == "ctransformers": MODEL_PATH = ensure_model_is_downloaded(self.llm_model_type) config = { "max_new_tokens": self.max_tokens_limit, "repetition_penalty": 1.1, } self.llm = CTransformers( model=MODEL_PATH, model_type="llama", config=config, callbacks=callbacks, ) elif self.llm_model_type == "hftgi": HFTGI_SERVER_URL = os.environ.get("HFTGI_SERVER_URL") self.max_tokens_limit = 4096 self.llm = HuggingFaceTextGenInference( inference_server_url=HFTGI_SERVER_URL, max_new_tokens=self.max_tokens_limit / 2, # top_k=0, top_p=0.95, # typical_p=0.95, temperature=0.01, repetition_penalty=1.115, callbacks=callbacks, timeout=600, streaming=True, ) elif self.llm_model_type.startswith("huggingface"): MODEL_NAME_OR_PATH = os.environ.get("HUGGINGFACE_MODEL_NAME_OR_PATH") print(f" loading model: {MODEL_NAME_OR_PATH}") hf_auth_token = ( os.environ.get("HUGGINGFACE_AUTH_TOKEN") if "Llama-2" in MODEL_NAME_OR_PATH else None ) transformers_offline = os.environ.get("TRANSFORMERS_OFFLINE") == "1" token = ( hf_auth_token if hf_auth_token is not None and len(hf_auth_token) > 0 and not transformers_offline else None ) print(f" HF auth token: {str(token)[-5:]}") if "Llama-2" in MODEL_NAME_OR_PATH: self.max_tokens_limit = 4096 is_t5 = "t5" in MODEL_NAME_OR_PATH temperature = ( 0.01 if "gpt4all-j" in MODEL_NAME_OR_PATH or "dolly" in MODEL_NAME_OR_PATH or "Qwen" in MODEL_NAME_OR_PATH or "Llama-2" in MODEL_NAME_OR_PATH else 0 ) use_fast = ( "stable" in MODEL_NAME_OR_PATH or "RedPajama" in MODEL_NAME_OR_PATH or "dolly" in MODEL_NAME_OR_PATH ) padding_side = "left" # if "dolly" in MODEL_NAME_OR_PATH else None config = ( AutoConfig.from_pretrained( MODEL_NAME_OR_PATH, trust_remote_code=True, token=token, fp32=hf_pipeline_device_type == "cpu", bf16=( hf_pipeline_device_type != "cpu" and torch_dtype == torch.bfloat16 ), fp16=( hf_pipeline_device_type != "cpu" and torch_dtype != torch.bfloat16 ), ) if "Qwen" in MODEL_NAME_OR_PATH else AutoConfig.from_pretrained( MODEL_NAME_OR_PATH, trust_remote_code=True, token=token, ) ) # config.attn_config["attn_impl"] = "triton" # config.max_seq_len = 4096 config.init_device = hf_pipeline_device_type tokenizer = ( T5Tokenizer.from_pretrained( MODEL_NAME_OR_PATH, token=token, ) if is_t5 else AutoTokenizer.from_pretrained( MODEL_NAME_OR_PATH, use_fast=use_fast, trust_remote_code=True, padding_side=padding_side, token=token, ) ) self._init_hf_streamer(tokenizer) task = "text2text-generation" if is_t5 else "text-generation" return_full_text = True if "dolly" in MODEL_NAME_OR_PATH else None repetition_penalty = ( 1.15 if "falcon" in MODEL_NAME_OR_PATH else (1.25 if "dolly" in MODEL_NAME_OR_PATH else 1.1) ) if load_quantized_model is not None: model = ( AutoModelForSeq2SeqLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, quantization_config=double_quant_config, trust_remote_code=True, use_auth_token=token, ) if is_t5 else AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, quantization_config=double_quant_config, trust_remote_code=True, use_auth_token=token, ) ) print(f"Model memory footprint: {model.get_memory_footprint()}") eos_token_id = -1 # starchat-beta uses a special <|end|> token with ID 49155 to denote ends of a turn if "starchat" in MODEL_NAME_OR_PATH: eos_token_id = 49155 pad_token_id = eos_token_id pipe = ( InstructionTextGenerationPipeline( task=task, model=model, tokenizer=tokenizer, streamer=self.streamer, max_new_tokens=2048, temperature=temperature, return_full_text=return_full_text, # langchain expects the full text repetition_penalty=repetition_penalty, ) if "dolly" in MODEL_NAME_OR_PATH else ( pipeline( task, model=model, tokenizer=tokenizer, eos_token_id=eos_token_id, pad_token_id=pad_token_id, streamer=self.streamer, return_full_text=return_full_text, # langchain expects the full text device_map="auto", trust_remote_code=True, max_new_tokens=2048, do_sample=True, temperature=0.01, top_p=0.95, top_k=50, repetition_penalty=repetition_penalty, ) if eos_token_id != -1 else pipeline( task, model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=return_full_text, # langchain expects the full text device_map="auto", trust_remote_code=True, max_new_tokens=2048, # verbose=True, temperature=temperature, top_p=0.95, top_k=0, # select from top 0 tokens (because zero, relies on top_p) repetition_penalty=repetition_penalty, ) ) ) elif "dolly" in MODEL_NAME_OR_PATH: model = AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, device_map=hf_pipeline_device_type, torch_dtype=torch_dtype, ) pipe = InstructionTextGenerationPipeline( task=task, model=model, tokenizer=tokenizer, streamer=self.streamer, max_new_tokens=2048, temperature=temperature, return_full_text=True, repetition_penalty=repetition_penalty, token=token, ) else: if os.environ.get("DISABLE_MODEL_PRELOADING") != "true": model = ( AutoModelForSeq2SeqLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, trust_remote_code=True, ) if is_t5 else ( AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, trust_remote_code=True, ) if "Qwen" in MODEL_NAME_OR_PATH else ( AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, trust_remote_code=True, ) if token is None else AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, trust_remote_code=True, token=token, ) ) ) ) print(f"Model memory footprint: {model.get_memory_footprint()}") model = model.eval() # print(f"Model memory footprint: {model.get_memory_footprint()}") else: model = MODEL_NAME_OR_PATH pipe = ( pipeline( task, model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=return_full_text, # langchain expects the full text device=hf_pipeline_device_type, torch_dtype=torch_dtype, max_new_tokens=2048, trust_remote_code=True, temperature=temperature, top_p=0.95, top_k=0, # select from top 0 tokens (because zero, relies on top_p) repetition_penalty=1.115, ) if token is None else pipeline( task, model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=return_full_text, # langchain expects the full text device=hf_pipeline_device_type, torch_dtype=torch_dtype, max_new_tokens=2048, trust_remote_code=True, temperature=temperature, top_p=0.95, top_k=0, # select from top 0 tokens (because zero, relies on top_p) repetition_penalty=1.115, token=token, ) ) self.llm = HuggingFacePipeline(pipeline=pipe, callbacks=callbacks) elif self.llm_model_type == "mosaicml": MODEL_NAME_OR_PATH = os.environ.get("MOSAICML_MODEL_NAME_OR_PATH") print(f" loading model: {MODEL_NAME_OR_PATH}") config = AutoConfig.from_pretrained( MODEL_NAME_OR_PATH, trust_remote_code=True ) # config.attn_config["attn_impl"] = "triton" config.max_seq_len = 16384 if "30b" in MODEL_NAME_OR_PATH else 4096 config.init_device = hf_pipeline_device_type model = ( AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, quantization_config=double_quant_config, trust_remote_code=True, ) if load_quantized_model is not None else AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, torch_dtype=torch_dtype, trust_remote_code=True, ) ) print(f"Model loaded on {config.init_device}") print(f"Model memory footprint: {model.get_memory_footprint()}") tokenizer = AutoTokenizer.from_pretrained("EleutherAI/gpt-neox-20b") self._init_hf_streamer(tokenizer) # mtp-7b is trained to add "<|endoftext|>" at the end of generations stop_token_ids = tokenizer.convert_tokens_to_ids(["<|endoftext|>"]) # define custom stopping criteria object class StopOnTokens(StoppingCriteria): def __call__( self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs, ) -> bool: for stop_id in stop_token_ids: if input_ids[0][-1] == stop_id: return True return False stopping_criteria = StoppingCriteriaList([StopOnTokens()]) max_new_tokens = 8192 if "30b" in MODEL_NAME_OR_PATH else 2048 self.max_tokens_limit = max_new_tokens self.search_kwargs = ( {"k": 8} if "30b" in MODEL_NAME_OR_PATH else self.search_kwargs ) repetition_penalty = 1.05 if "30b" in MODEL_NAME_OR_PATH else 1.02 pipe = ( pipeline( model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=True, # langchain expects the full text task="text-generation", device_map="auto", # we pass model parameters here too stopping_criteria=stopping_criteria, # without this model will ramble temperature=0, # 'randomness' of outputs, 0.0 is the min and 1.0 the max top_p=0.95, # select from top tokens whose probability add up to 15% top_k=0, # select from top 0 tokens (because zero, relies on top_p) max_new_tokens=max_new_tokens, # mex number of tokens to generate in the output repetition_penalty=repetition_penalty, # without this output begins repeating ) if load_quantized_model is not None else pipeline( model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=True, # langchain expects the full text task="text-generation", device=config.init_device, # we pass model parameters here too stopping_criteria=stopping_criteria, # without this model will ramble temperature=0, # 'randomness' of outputs, 0.0 is the min and 1.0 the max top_p=0.95, # select from top tokens whose probability add up to 15% top_k=0, # select from top 0 tokens (because zero, relies on top_p) max_new_tokens=max_new_tokens, # mex number of tokens to generate in the output repetition_penalty=repetition_penalty, # without this output begins repeating ) ) self.llm = HuggingFacePipeline(pipeline=pipe, callbacks=callbacks) elif self.llm_model_type == "stablelm": MODEL_NAME_OR_PATH = os.environ.get("STABLELM_MODEL_NAME_OR_PATH") print(f" loading model: {MODEL_NAME_OR_PATH}") config = AutoConfig.from_pretrained( MODEL_NAME_OR_PATH, trust_remote_code=True ) # config.attn_config["attn_impl"] = "triton" # config.max_seq_len = 4096 config.init_device = hf_pipeline_device_type model = ( AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, quantization_config=double_quant_config, trust_remote_code=True, ) if load_quantized_model is not None else AutoModelForCausalLM.from_pretrained( MODEL_NAME_OR_PATH, config=config, torch_dtype=torch_dtype, trust_remote_code=True, ) ) print(f"Model loaded on {config.init_device}") print(f"Model memory footprint: {model.get_memory_footprint()}") tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME_OR_PATH) self._init_hf_streamer(tokenizer) class StopOnTokens(StoppingCriteria): def __call__( self, input_ids: torch.LongTensor, scores: torch.FloatTensor, **kwargs, ) -> bool: stop_ids = [50278, 50279, 50277, 1, 0] for stop_id in stop_ids: if input_ids[0][-1] == stop_id: return True return False stopping_criteria = StoppingCriteriaList([StopOnTokens()]) pipe = ( pipeline( model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=True, # langchain expects the full text task="text-generation", device_map="auto", # we pass model parameters here too stopping_criteria=stopping_criteria, # without this model will ramble temperature=0, # 'randomness' of outputs, 0.0 is the min and 1.0 the max top_p=0.95, # select from top tokens whose probability add up to 15% top_k=0, # select from top 0 tokens (because zero, relies on top_p) max_new_tokens=2048, # mex number of tokens to generate in the output repetition_penalty=1.25, # without this output begins repeating ) if load_quantized_model is not None else pipeline( model=model, tokenizer=tokenizer, streamer=self.streamer, return_full_text=True, # langchain expects the full text task="text-generation", device=config.init_device, # we pass model parameters here too stopping_criteria=stopping_criteria, # without this model will ramble temperature=0, # 'randomness' of outputs, 0.0 is the min and 1.0 the max top_p=0.95, # select from top tokens whose probability add up to 15% top_k=0, # select from top 0 tokens (because zero, relies on top_p) max_new_tokens=2048, # mex number of tokens to generate in the output repetition_penalty=1.05, # without this output begins repeating ) ) self.llm = HuggingFacePipeline(pipeline=pipe, callbacks=callbacks) print("initialization complete")