Spaces:

iblfe
/

test

Runtime error

App Files Files Community

test / src /llm_exllama.py

iblfe

Upload folder using huggingface_hub

b585c7f verified 5 months ago

raw history blame contribute delete

No virus

14.7 kB

	from functools import partial

	from langchain.llms.base import LLM
	from langchain.callbacks.manager import CallbackManagerForLLMRun
	from typing import Any, Dict, List, Optional
	from exllama.model import ExLlama, ExLlamaCache, ExLlamaConfig
	from exllama.tokenizer import ExLlamaTokenizer
	from exllama.generator import ExLlamaGenerator
	from exllama.lora import ExLlamaLora
	import os, glob

	from pydantic.v1 import root_validator

	BROKEN_UNICODE = b'\\ufffd'.decode('unicode_escape')

	class H2OExLlamaTokenizer(ExLlamaTokenizer):
	def __call__(self, text, args, *kwargs):
	return dict(input_ids=self.encode(text))


	class H2OExLlamaGenerator(ExLlamaGenerator):
	def is_exlama(self):
	return True


	class Exllama(LLM):
	client: Any #: :meta private:
	model_path: str = None
	model: Any = None
	sanitize_bot_response: bool = False
	prompter: Any = None
	context: Any = ''
	iinput: Any = ''

	"""The path to the GPTQ model folder."""
	exllama_cache: ExLlamaCache = None #: :meta private:
	config: ExLlamaConfig = None #: :meta private:
	generator: ExLlamaGenerator = None #: :meta private:
	tokenizer: ExLlamaTokenizer = None #: :meta private:

	##Langchain parameters
	logfunc = print
	stop_sequences: Optional[List[str]] = "" #, description="Sequences that immediately will stop the generator.")
	streaming: Optional[bool] = True #, description="Whether to stream the results, token by token.")

	##Generator parameters
	disallowed_tokens: Optional[List[int]] = None # description="List of tokens to disallow during generation.")
	temperature: Optional[float] = None # description="Temperature for sampling diversity.")
	top_k: Optional[int] = None # description="Consider the most probable top_k samples, 0 to disable top_k sampling.")
	top_p: Optional[float] = None # description="Consider tokens up to a cumulative probabiltiy of top_p, 0.0 to disable top_p sampling.")
	min_p: Optional[float] = None # description="Do not consider tokens with probability less than this.")
	typical: Optional[float] = None # description="Locally typical sampling threshold, 0.0 to disable typical sampling.")
	token_repetition_penalty_max: Optional[float] = None # description="Repetition penalty for most recent tokens.")
	token_repetition_penalty_sustain: Optional[int] = None # description="No. most recent tokens to repeat penalty for, -1 to apply to whole context.")
	token_repetition_penalty_decay: Optional[int] = None # description="Gradually decrease penalty over this many tokens.")
	beams: Optional[int] = None # description="Number of beams for beam search.")
	beam_length: Optional[int] = None # description="Length of beams for beam search.")

	##Config overrides
	max_seq_len: Optional[int] = 2048 # decription="Reduce to save memory. Can also be increased, ideally while also using compress_pos_emn and a compatible model/LoRA")
	compress_pos_emb: Optional[float] = 1.0 # description="Amount of compression to apply to the positional embedding.")
	set_auto_map: Optional[str] = None # description="Comma-separated list of VRAM (in GB) to use per GPU device for model layers, e.g. 20,7,7")
	gpu_peer_fix: Optional[bool] = None # description="Prevent direct copies of data between GPUs")
	alpha_value: Optional[float] = 1.0 #, description="Rope context extension alpha")

	##Tuning
	matmul_recons_thd: Optional[int] = None
	fused_mlp_thd: Optional[int] = None
	sdp_thd: Optional[int] = None
	fused_attn: Optional[bool] = None
	matmul_fused_remap: Optional[bool] = None
	rmsnorm_no_half2: Optional[bool] = None
	rope_no_half2: Optional[bool] = None
	matmul_no_half2: Optional[bool] = None
	silu_no_half2: Optional[bool] = None
	concurrent_streams: Optional[bool] = None

	##Lora Parameters
	lora_path: Optional[str] = None # description="Path to your lora.")

	@staticmethod
	def get_model_path_at(path):
	patterns = [".safetensors", ".bin", "*.pt"]
	model_paths = []
	for pattern in patterns:
	full_pattern = os.path.join(path, pattern)
	model_paths = glob.glob(full_pattern)
	if model_paths: # If there are any files matching the current pattern
	break # Exit the loop as soon as we find a matching file
	if model_paths: # If there are any files matching any of the patterns
	return model_paths[0]
	else:
	return None # Return None if no matching files were found

	@staticmethod
	def configure_object(params, values, logfunc):
	obj_params = {k: values.get(k) for k in params}

	def apply_to(obj):
	for key, value in obj_params.items():
	if value:
	if hasattr(obj, key):
	setattr(obj, key, value)
	logfunc(f"{key} {value}")
	else:
	raise AttributeError(f"{key} does not exist in {obj}")

	return apply_to

	@root_validator()
	def validate_environment(cls, values: Dict) -> Dict:
	model_param_names = [
	"temperature",
	"top_k",
	"top_p",
	"min_p",
	"typical",
	"token_repetition_penalty_max",
	"token_repetition_penalty_sustain",
	"token_repetition_penalty_decay",
	"beams",
	"beam_length",
	]

	config_param_names = [
	"max_seq_len",
	"compress_pos_emb",
	"gpu_peer_fix",
	"alpha_value"
	]

	tuning_parameters = [
	"matmul_recons_thd",
	"fused_mlp_thd",
	"sdp_thd",
	"matmul_fused_remap",
	"rmsnorm_no_half2",
	"rope_no_half2",
	"matmul_no_half2",
	"silu_no_half2",
	"concurrent_streams",
	"fused_attn",
	]

	##Set logging function if verbose or set to empty lambda
	verbose = values['verbose']
	if not verbose:
	values['logfunc'] = lambda args, *kwargs: None
	logfunc = values['logfunc']

	if values['model'] is None:
	model_path = values["model_path"]
	lora_path = values["lora_path"]

	tokenizer_path = os.path.join(model_path, "tokenizer.model")
	model_config_path = os.path.join(model_path, "config.json")
	model_path = Exllama.get_model_path_at(model_path)

	config = ExLlamaConfig(model_config_path)
	tokenizer = ExLlamaTokenizer(tokenizer_path)
	config.model_path = model_path

	configure_config = Exllama.configure_object(config_param_names, values, logfunc)
	configure_config(config)
	configure_tuning = Exllama.configure_object(tuning_parameters, values, logfunc)
	configure_tuning(config)

	##Special parameter, set auto map, it's a function
	if values['set_auto_map']:
	config.set_auto_map(values['set_auto_map'])
	logfunc(f"set_auto_map {values['set_auto_map']}")

	model = ExLlama(config)
	exllama_cache = ExLlamaCache(model)
	generator = ExLlamaGenerator(model, tokenizer, exllama_cache)

	##Load and apply lora to generator
	if lora_path is not None:
	lora_config_path = os.path.join(lora_path, "adapter_config.json")
	lora_path = Exllama.get_model_path_at(lora_path)
	lora = ExLlamaLora(model, lora_config_path, lora_path)
	generator.lora = lora
	logfunc(f"Loaded LORA @ {lora_path}")
	else:
	generator = values['model']
	exllama_cache = generator.cache
	model = generator.model
	config = model.config
	tokenizer = generator.tokenizer

	# Set if model existed before or not since generation-time parameters
	configure_model = Exllama.configure_object(model_param_names, values, logfunc)
	values["stop_sequences"] = [x.strip().lower() for x in values["stop_sequences"]]
	configure_model(generator.settings)

	setattr(generator.settings, "stop_sequences", values["stop_sequences"])
	logfunc(f"stop_sequences {values['stop_sequences']}")

	disallowed = values.get("disallowed_tokens")
	if disallowed:
	generator.disallow_tokens(disallowed)
	print(f"Disallowed Tokens: {generator.disallowed_tokens}")

	values["client"] = model
	values["generator"] = generator
	values["config"] = config
	values["tokenizer"] = tokenizer
	values["exllama_cache"] = exllama_cache

	return values

	@property
	def _llm_type(self) -> str:
	"""Return type of llm."""
	return "Exllama"

	def get_num_tokens(self, text: str) -> int:
	"""Get the number of tokens present in the text."""
	return self.generator.tokenizer.num_tokens(text)

	def get_token_ids(self, text: str) -> List[int]:
	return self.generator.tokenizer.encode(text)
	# avoid base method that is not aware of how to properly tokenize (uses GPT2)
	# return _get_token_ids_default_method(text)

	def _call(
	self,
	prompt: str,
	stop: Optional[List[str]] = None,
	run_manager: Optional[CallbackManagerForLLMRun] = None,
	**kwargs: Any,
	) -> str:
	assert self.tokenizer is not None
	from h2oai_pipeline import H2OTextGenerationPipeline
	prompt, num_prompt_tokens = H2OTextGenerationPipeline.limit_prompt(prompt, self.tokenizer)

	# NOTE: TGI server does not add prompting, so must do here
	data_point = dict(context=self.context, instruction=prompt, input=self.iinput)
	prompt = self.prompter.generate_prompt(data_point)

	text = ''
	for text1 in self.stream(prompt=prompt, stop=stop, run_manager=run_manager):
	text = text1
	return text

	from enum import Enum

	class MatchStatus(Enum):
	EXACT_MATCH = 1
	PARTIAL_MATCH = 0
	NO_MATCH = 2

	def match_status(self, sequence: str, banned_sequences: List[str]):
	sequence = sequence.strip().lower()
	for banned_seq in banned_sequences:
	if banned_seq == sequence:
	return self.MatchStatus.EXACT_MATCH
	elif banned_seq.startswith(sequence):
	return self.MatchStatus.PARTIAL_MATCH
	return self.MatchStatus.NO_MATCH

	def stream(
	self,
	prompt: str,
	stop: Optional[List[str]] = None,
	run_manager: Optional[CallbackManagerForLLMRun] = None,
	) -> str:
	config = self.config
	generator = self.generator
	beam_search = (self.beams and self.beams >= 1 and self.beam_length and self.beam_length >= 1)

	ids = generator.tokenizer.encode(prompt)
	generator.gen_begin_reuse(ids)

	if beam_search:
	generator.begin_beam_search()
	token_getter = generator.beam_search
	else:
	generator.end_beam_search()
	token_getter = generator.gen_single_token

	last_newline_pos = 0
	seq_length = len(generator.tokenizer.decode(generator.sequence_actual[0]))
	response_start = seq_length
	cursor_head = response_start

	text_callback = None
	if run_manager:
	text_callback = partial(
	run_manager.on_llm_new_token, verbose=self.verbose
	)
	# No longer assume below, assume always just new text so various langchain things work
	##### parent handler of streamer expects to see prompt first else output="" and lose if prompt=None in prompter
	#### text_callback:
	#### text_callback(prompt)
	text = ""
	while (generator.gen_num_tokens() <= (
	self.max_seq_len - 4)): # Slight extra padding space as we seem to occassionally get a few more than 1-2 tokens
	# Fetch a token
	token = token_getter()

	# If it's the ending token replace it and end the generation.
	if token.item() == generator.tokenizer.eos_token_id:
	generator.replace_last_token(generator.tokenizer.newline_token_id)
	if beam_search:
	generator.end_beam_search()
	return

	# Tokenize the string from the last new line, we can't just decode the last token due to how sentencepiece decodes.
	stuff = generator.tokenizer.decode(generator.sequence_actual[0][last_newline_pos:])
	cursor_tail = len(stuff)
	has_unicode_combined = cursor_tail<cursor_head
	text_chunk = stuff[cursor_head:cursor_tail]
	if has_unicode_combined:
	# replace the broken unicode character with combined one
	text=text[:-2]
	text_chunk = stuff[cursor_tail-1:cursor_tail]

	cursor_head = cursor_tail

	# Append the generated chunk to our stream buffer
	text += text_chunk
	text = self.prompter.get_response(prompt + text, prompt=prompt,
	sanitize_bot_response=self.sanitize_bot_response)

	if token.item() == generator.tokenizer.newline_token_id:
	last_newline_pos = len(generator.sequence_actual[0])
	cursor_head = 0
	cursor_tail = 0

	# Check if the stream buffer is one of the stop sequences
	status = self.match_status(text, self.stop_sequences)

	if status == self.MatchStatus.EXACT_MATCH:
	# Encountered a stop, rewind our generator to before we hit the match and end generation.
	rewind_length = generator.tokenizer.encode(text).shape[-1]
	generator.gen_rewind(rewind_length)
	#gen = generator.tokenizer.decode(generator.sequence_actual[0][response_start:])
	if beam_search:
	generator.end_beam_search()
	return
	elif status == self.MatchStatus.PARTIAL_MATCH:
	# Partially matched a stop, continue buffering but don't yield.
	continue
	elif status == self.MatchStatus.NO_MATCH:
	if text_callback and not (text_chunk == BROKEN_UNICODE):
	text_callback(text_chunk)
	yield text # Not a stop, yield the match buffer.

	return