llama2_wrapper/model.py

import os
import time
import uuid
from enum import Enum
from threading import Thread
from typing import Any, Iterator, Union, List
from llama2_wrapper.types import (
 Completion,
 CompletionChunk,
 ChatCompletion,
 ChatCompletionChunk,
 # ChatCompletionMessage,
 Message,
 B_INST,
 E_INST,
 B_SYS,
 E_SYS,
)


class LLAMA2_WRAPPER:
 def __init__(
 self,
 model_path: str = "",
 backend_type: str = "llama.cpp",
 max_tokens: int = 4000,
 load_in_8bit: bool = True,
 verbose: bool = False,
 ):
 """Load a llama2 model from `model_path`.

 Args:
 model_path: Path to the model.
 backend_type: Backend for llama2, options: llama.cpp, gptq, transformers
 max_tokens: Maximum context size.
 load_in_8bit: Use bitsandbytes to run model in 8 bit mode (only for transformers models).
 verbose: Print verbose output to stderr.

 Raises:
 ValueError: If the model path does not exist.

 Returns:
 A LLAMA2_WRAPPER instance.
 """
 self.model_path = model_path
 self.backend_type = BackendType.get_type(backend_type)
 self.max_tokens = max_tokens
 self.load_in_8bit = load_in_8bit

 self.model = None
 self.tokenizer = None

 self.verbose = verbose

 if self.backend_type is BackendType.LLAMA_CPP:
 print("Running on backend llama.cpp.")
 else:
 import torch

 if torch.cuda.is_available():
 print("Running on GPU with backend torch transformers.")
 else:
 print("GPU CUDA not found.")

 self.default_llamacpp_path = "./models/llama-2-7b-chat.Q4_0.gguf"
 self.default_gptq_path = "./models/Llama-2-7b-Chat-GPTQ"
 # Download default ggml/gptq model
 if self.model_path == "":
 print("Model path is empty.")
 if self.backend_type is BackendType.LLAMA_CPP:
 print("Use default llama.cpp model path: " + self.default_llamacpp_path)
 if not os.path.exists(self.default_llamacpp_path):
 print("Start downloading model to: " + self.default_llamacpp_path)
 from huggingface_hub import hf_hub_download

 hf_hub_download(
 repo_id="TheBloke/Llama-2-7b-Chat-GGUF",
 filename="llama-2-7b-chat.Q4_0.gguf",
 local_dir="./models/",
 )
 else:
 print("Model exists in ./models/llama-2-7b-chat.Q4_0.gguf.")
 self.model_path = self.default_llamacpp_path
 elif self.backend_type is BackendType.GPTQ:
 print("Use default gptq model path: " + self.default_gptq_path)
 if not os.path.exists(self.default_gptq_path):
 print("Start downloading model to: " + self.default_gptq_path)
 from huggingface_hub import snapshot_download

 snapshot_download(
 "TheBloke/Llama-2-7b-Chat-GPTQ",
 local_dir=self.default_gptq_path,
 )
 else:
 print("Model exists in " + self.default_gptq_path)
 self.model_path = self.default_gptq_path

 self.init_tokenizer()
 self.init_model()

 def init_model(self):
 if self.model is None:
 self.model = LLAMA2_WRAPPER.create_llama2_model(
 self.model_path,
 self.backend_type,
 self.max_tokens,
 self.load_in_8bit,
 self.verbose,
 )
 if self.backend_type is not BackendType.LLAMA_CPP:
 self.model.eval()

 def init_tokenizer(self):
 if self.backend_type is not BackendType.LLAMA_CPP:
 if self.tokenizer is None:
 self.tokenizer = LLAMA2_WRAPPER.create_llama2_tokenizer(self.model_path)

 @classmethod
 def create_llama2_model(
 cls, model_path, backend_type, max_tokens, load_in_8bit, verbose
 ):
 if backend_type is BackendType.LLAMA_CPP:
 from llama_cpp import Llama

 model = Llama(
 model_path=model_path,
 n_ctx=max_tokens,
 n_batch=max_tokens,
 verbose=verbose,
 )
 elif backend_type is BackendType.GPTQ:
 from auto_gptq import AutoGPTQForCausalLM

 model = AutoGPTQForCausalLM.from_quantized(
 model_path,
 use_safetensors=True,
 trust_remote_code=True,
 device="cuda:0",
 use_triton=False,
 quantize_config=None,
 )
 elif backend_type is BackendType.TRANSFORMERS:
 import torch
 from transformers import AutoModelForCausalLM

 model = AutoModelForCausalLM.from_pretrained(
 model_path,
 device_map="auto",
 torch_dtype=torch.float16,
 load_in_8bit=load_in_8bit,
 )
 else:
 print(backend_type + "not implemented.")
 return model

 @classmethod
 def create_llama2_tokenizer(cls, model_path):
 from transformers import AutoTokenizer

 tokenizer = AutoTokenizer.from_pretrained(model_path)
 return tokenizer

 def get_token_length(
 self,
 prompt: str,
 ) -> int:
 if self.backend_type is BackendType.LLAMA_CPP:
 input_ids = self.model.tokenize(bytes(prompt, "utf-8"))
 return len(input_ids)
 else:
 input_ids = self.tokenizer([prompt], return_tensors="np")["input_ids"]
 return input_ids.shape[-1]

 def get_input_token_length(
 self,
 message: str,
 chat_history: list[tuple[str, str]] = [],
 system_prompt: str = "",
 ) -> int:
 prompt = get_prompt(message, chat_history, system_prompt)

 return self.get_token_length(prompt)

 def generate(
 self,
 prompt: str,
 max_new_tokens: int = 1000,
 temperature: float = 0.9,
 top_p: float = 1.0,
 top_k: int = 40,
 repetition_penalty: float = 1.0,
 **kwargs: Any,
 ) -> Iterator[str]:
 """Create a generator of response from a prompt.

 Examples:
 >>> llama2_wrapper = LLAMA2_WRAPPER()
 >>> prompt = get_prompt("Hi do you know Pytorch?")
 >>> for response in llama2_wrapper.generate(prompt):
 ... print(response)

 Args:
 prompt: The prompt to generate text from.
 max_new_tokens: The maximum number of tokens to generate.
 temperature: The temperature to use for sampling.
 top_p: The top-p value to use for sampling.
 top_k: The top-k value to use for sampling.
 repetition_penalty: The penalty to apply to repeated tokens.
 kwargs: all other arguments.

 Yields:
 The generated text.
 """
 if self.backend_type is BackendType.LLAMA_CPP:
 result = self.model(
 prompt=prompt,
 stream=True,
 max_tokens=max_new_tokens,
 top_k=top_k,
 top_p=top_p,
 temperature=temperature,
 repeat_penalty=repetition_penalty,
 **kwargs,
 )
 outputs = []
 for part in result:
 text = part["choices"][0]["text"]
 outputs.append(text)
 yield "".join(outputs)
 else:
 from transformers import TextIteratorStreamer

 inputs = self.tokenizer([prompt], return_tensors="pt").to("cuda")

 streamer = TextIteratorStreamer(
 self.tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
 )
 generate_kwargs = dict(
 inputs,
 streamer=streamer,
 max_new_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repetition_penalty=repetition_penalty,
 # num_beams=1,
 )
 generate_kwargs = (
 generate_kwargs if kwargs is None else {**generate_kwargs, **kwargs}
 )
 t = Thread(target=self.model.generate, kwargs=generate_kwargs)
 t.start()

 outputs = []
 for text in streamer:
 outputs.append(text)
 yield "".join(outputs)

 def run(
 self,
 message: str,
 chat_history: list[tuple[str, str]] = [],
 system_prompt: str = "",
 max_new_tokens: int = 1000,
 temperature: float = 0.9,
 top_p: float = 1.0,
 top_k: int = 40,
 repetition_penalty: float = 1.0,
 ) -> Iterator[str]:
 """Create a generator of response from a chat message.
 Process message to llama2 prompt with chat history
 and system_prompt for chatbot.

 Args:
 message: The origianl chat message to generate text from.
 chat_history: Chat history list from chatbot.
 system_prompt: System prompt for chatbot.
 max_new_tokens: The maximum number of tokens to generate.
 temperature: The temperature to use for sampling.
 top_p: The top-p value to use for sampling.
 top_k: The top-k value to use for sampling.
 repetition_penalty: The penalty to apply to repeated tokens.
 kwargs: all other arguments.

 Yields:
 The generated text.
 """
 prompt = get_prompt(message, chat_history, system_prompt)
 return self.generate(
 prompt, max_new_tokens, temperature, top_p, top_k, repetition_penalty
 )

 def __call__(
 self,
 prompt: str,
 stream: bool = False,
 max_new_tokens: int = 1000,
 temperature: float = 0.9,
 top_p: float = 1.0,
 top_k: int = 40,
 repetition_penalty: float = 1.0,
 **kwargs: Any,
 ) -> Union[str, Iterator[str]]:
 """Generate text from a prompt.

 Examples:
 >>> llama2_wrapper = LLAMA2_WRAPPER()
 >>> prompt = get_prompt("Hi do you know Pytorch?")
 >>> print(llama2_wrapper(prompt))

 Args:
 prompt: The prompt to generate text from.
 stream: Whether to stream the results.
 max_new_tokens: The maximum number of tokens to generate.
 temperature: The temperature to use for sampling.
 top_p: The top-p value to use for sampling.
 top_k: The top-k value to use for sampling.
 repetition_penalty: The penalty to apply to repeated tokens.
 kwargs: all other arguments.

 Raises:
 ValueError: If the requested tokens exceed the context window.
 RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

 Returns:
 Generated text.
 """
 if self.backend_type is BackendType.LLAMA_CPP:
 completion_or_chunks = self.model.__call__(
 prompt,
 stream=stream,
 max_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repeat_penalty=repetition_penalty,
 **kwargs,
 )
 if stream:

 def chunk_generator(chunks):
 for part in chunks:
 chunk = part["choices"][0]["text"]
 yield chunk

 chunks: Iterator[str] = chunk_generator(completion_or_chunks)
 return chunks
 return completion_or_chunks["choices"][0]["text"]
 else:
 inputs = self.tokenizer([prompt], return_tensors="pt").input_ids
 prompt_tokens_len = len(inputs[0])
 inputs = inputs.to("cuda")
 generate_kwargs = dict(
 inputs=inputs,
 max_new_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repetition_penalty=repetition_penalty,
 # num_beams=1,
 )
 generate_kwargs = (
 generate_kwargs if kwargs is None else {**generate_kwargs, **kwargs}
 )
 if stream:
 from transformers import TextIteratorStreamer

 streamer = TextIteratorStreamer(
 self.tokenizer,
 timeout=10.0,
 skip_prompt=True,
 skip_special_tokens=True,
 )
 generate_kwargs["streamer"] = streamer

 t = Thread(target=self.model.generate, kwargs=generate_kwargs)
 t.start()
 return streamer
 else:
 output_ids = self.model.generate(
 **generate_kwargs,
 )
 # skip prompt, skip special tokens
 output = self.tokenizer.decode(
 output_ids[0][prompt_tokens_len:], skip_special_tokens=True
 )
 return output

 def completion(
 self,
 prompt: str,
 stream: bool = False,
 max_new_tokens: int = 1000,
 temperature: float = 0.9,
 top_p: float = 1.0,
 top_k: int = 40,
 repetition_penalty: float = 1.0,
 **kwargs: Any,
 ) -> Union[Completion, Iterator[CompletionChunk]]:
 """For OpenAI compatible API /v1/completions
 Generate text from a prompt.

 Examples:
 >>> llama2_wrapper = LLAMA2_WRAPPER()
 >>> prompt = get_prompt("Hi do you know Pytorch?")
 >>> print(llm.completion(prompt))

 Args:
 prompt: The prompt to generate text from.
 stream: Whether to stream the results.
 max_new_tokens: The maximum number of tokens to generate.
 temperature: The temperature to use for sampling.
 top_p: The top-p value to use for sampling.
 top_k: The top-k value to use for sampling.
 repetition_penalty: The penalty to apply to repeated tokens.
 kwargs: all other arguments.

 Raises:
 ValueError: If the requested tokens exceed the context window.
 RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

 Returns:
 Response object containing the generated text.
 """
 completion_id: str = f"cmpl-{str(uuid.uuid4())}"
 created: int = int(time.time())
 model_name: str = (
 self.backend_type + " default model"
 if self.model_path == ""
 else self.model_path
 )
 if self.backend_type is BackendType.LLAMA_CPP:
 completion_or_chunks = self.model.__call__(
 prompt,
 stream=stream,
 max_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repeat_penalty=repetition_penalty,
 **kwargs,
 )
 if stream:
 chunks: Iterator[CompletionChunk] = completion_or_chunks
 return chunks
 return completion_or_chunks
 else:
 inputs = self.tokenizer([prompt], return_tensors="pt").input_ids
 prompt_tokens_len = len(inputs[0])
 inputs = inputs.to("cuda")
 generate_kwargs = dict(
 inputs=inputs,
 max_new_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repetition_penalty=repetition_penalty,
 # num_beams=1,
 )
 generate_kwargs = (
 generate_kwargs if kwargs is None else {**generate_kwargs, **kwargs}
 )
 if stream:
 from transformers import TextIteratorStreamer

 streamer = TextIteratorStreamer(
 self.tokenizer,
 timeout=10.0,
 skip_prompt=True,
 skip_special_tokens=True,
 )
 generate_kwargs["streamer"] = streamer

 t = Thread(target=self.model.generate, kwargs=generate_kwargs)
 t.start()

 def chunk_generator(chunks):
 for part in chunks:
 yield {
 "id": completion_id,
 "object": "text_completion",
 "created": created,
 "model": model_name,
 "choices": [
 {
 "text": part,
 "index": 0,
 "logprobs": None,
 "finish_reason": None,
 }
 ],
 }

 chunks: Iterator[CompletionChunk] = chunk_generator(streamer)
 return chunks

 else:
 output_ids = self.model.generate(
 **generate_kwargs,
 )
 total_tokens_len = len(output_ids[0])
 output = self.tokenizer.decode(
 output_ids[0][prompt_tokens_len:], skip_special_tokens=True
 )
 completion: Completion = {
 "id": completion_id,
 "object": "text_completion",
 "created": created,
 "model": model_name,
 "choices": [
 {
 "text": output,
 "index": 0,
 "logprobs": None,
 "finish_reason": None,
 }
 ],
 "usage": {
 "prompt_tokens": prompt_tokens_len,
 "completion_tokens": total_tokens_len - prompt_tokens_len,
 "total_tokens": total_tokens_len,
 },
 }
 return completion

 def chat_completion(
 self,
 messages: List[Message],
 stream: bool = False,
 max_new_tokens: int = 1000,
 temperature: float = 0.9,
 top_p: float = 1.0,
 top_k: int = 40,
 repetition_penalty: float = 1.0,
 **kwargs: Any,
 ) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
 """For OpenAI compatible API /v1/chat/completions
 Generate text from a dialog (chat history).

 Examples:
 >>> llama2_wrapper = LLAMA2_WRAPPER()
 >>> dialog = [
 {
 "role":"system",
 "content":"You are a helpful, respectful and honest assistant. "
 },{
 "role":"user",
 "content":"Hi do you know Pytorch?",
 },
 ]
 >>> print(llm.chat_completion(dialog))

 Args:
 dialog: The dialog (chat history) to generate text from.
 stream: Whether to stream the results.
 max_new_tokens: The maximum number of tokens to generate.
 temperature: The temperature to use for sampling.
 top_p: The top-p value to use for sampling.
 top_k: The top-k value to use for sampling.
 repetition_penalty: The penalty to apply to repeated tokens.
 kwargs: all other arguments.

 Raises:
 ValueError: If the requested tokens exceed the context window.
 RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.

 Returns:
 Response object containing the generated text.
 """
 completion_id: str = f"cmpl-{str(uuid.uuid4())}"
 created: int = int(time.time())
 model_name: str = (
 self.backend_type + " default model"
 if self.model_path == ""
 else self.model_path
 )
 if self.backend_type is BackendType.LLAMA_CPP:
 completion_or_chunks = self.model.create_chat_completion(
 messages,
 stream=stream,
 max_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repeat_penalty=repetition_penalty,
 **kwargs,
 )
 if stream:
 chunks: Iterator[ChatCompletionChunk] = completion_or_chunks
 return chunks
 return completion_or_chunks
 else:
 prompt = get_prompt_for_dialog(messages)
 inputs = self.tokenizer([prompt], return_tensors="pt").input_ids
 prompt_tokens_len = len(inputs[0])
 inputs = inputs.to("cuda")
 generate_kwargs = dict(
 inputs=inputs,
 max_new_tokens=max_new_tokens,
 temperature=temperature,
 top_p=top_p,
 top_k=top_k,
 repetition_penalty=repetition_penalty,
 # num_beams=1,
 )
 generate_kwargs = (
 generate_kwargs if kwargs is None else {**generate_kwargs, **kwargs}
 )
 if stream:
 from transformers import TextIteratorStreamer

 streamer = TextIteratorStreamer(
 self.tokenizer,
 timeout=10.0,
 skip_prompt=True,
 skip_special_tokens=True,
 )
 generate_kwargs["streamer"] = streamer
 t = Thread(target=self.model.generate, kwargs=generate_kwargs)
 t.start()

 def chunk_generator(chunks):
 yield {
 "id": "chat" + completion_id,
 "model": model_name,
 "created": created,
 "object": "chat.completion.chunk",
 "choices": [
 {
 "index": 0,
 "delta": {
 "role": "assistant",
 },
 "finish_reason": None,
 }
 ],
 }
 for part in enumerate(chunks):
 yield {
 "id": "chat" + completion_id,
 "model": model_name,
 "created": created,
 "object": "chat.completion.chunk",
 "choices": [
 {
 "index": 0,
 "delta": {
 "content": part,
 },
 "finish_reason": None,
 }
 ],
 }

 chunks: Iterator[ChatCompletionChunk] = chunk_generator(streamer)
 return chunks

 else:
 output_ids = self.model.generate(
 **generate_kwargs,
 )
 total_tokens_len = len(output_ids[0])
 output = self.tokenizer.decode(
 output_ids[0][prompt_tokens_len:], skip_special_tokens=True
 )
 chatcompletion: ChatCompletion = {
 "id": "chat" + completion_id,
 "object": "chat.completion",
 "created": created,
 "model": model_name,
 "choices": [
 {
 "index": 0,
 "message": {
 "role": "assistant",
 "content": output,
 },
 "finish_reason": None,
 }
 ],
 "usage": {
 "prompt_tokens": prompt_tokens_len,
 "completion_tokens": total_tokens_len - prompt_tokens_len,
 "total_tokens": total_tokens_len,
 },
 }
 return chatcompletion


def get_prompt_for_dialog(dialog: List[Message]) -> str:
 """Process dialog (chat history) to llama2 prompt for
 OpenAI compatible API /v1/chat/completions.

 Examples:
 >>> dialog = [
 {
 "role":"system",
 "content":"You are a helpful, respectful and honest assistant. "
 },{
 "role":"user",
 "content":"Hi do you know Pytorch?",
 },
 ]
 >>> prompt = get_prompt_for_dialog("Hi do you know Pytorch?")

 Args:
 dialog: The dialog (chat history) to generate text from.

 Yields:
 prompt string.
 """
 # add "<<SYS>>\n{system_prompt}\n<</SYS>>\n\n" in first dialog
 if dialog[0]["role"] == "system":
 dialog = [
 {
 "role": dialog[1]["role"],
 "content": B_SYS + dialog[0]["content"] + E_SYS + dialog[1]["content"],
 }
 ] + dialog[2:]
 # check roles
 assert all([msg["role"] == "user" for msg in dialog[::2]]) and all(
 [msg["role"] == "assistant" for msg in dialog[1::2]]
 ), (
 "model only supports 'system', 'user' and 'assistant' roles, "
 "starting with 'system', then 'user' and alternating (u/a/u/a/u...)"
 )
 # add chat history
 texts = []
 for prompt, answer in zip(
 dialog[::2],
 dialog[1::2],
 ):
 texts.append(
 f"{B_INST} {(prompt['content']).strip()} {E_INST} {(answer['content']).strip()} "
 )
 # check last message if role is user, then add it to prompt text
 assert (
 dialog[-1]["role"] == "user"
 ), f"Last message must be from user, got {dialog[-1]['role']}"
 texts.append(f"{B_INST} {(dialog[-1]['content']).strip()} {E_INST}")
 return "".join(texts)


def get_prompt(
 message: str, chat_history: list[tuple[str, str]] = [], system_prompt: str = ""
) -> str:
 """Process message to llama2 prompt with chat history
 and system_prompt for chatbot.

 Examples:
 >>> prompt = get_prompt("Hi do you know Pytorch?")

 Args:
 message: The origianl chat message to generate text from.
 chat_history: Chat history list from chatbot.
 system_prompt: System prompt for chatbot.

 Yields:
 prompt string.
 """
 texts = [f"[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n\n"]
 for user_input, response in chat_history:
 texts.append(f"{user_input.strip()} [/INST] {response.strip()} </s><s> [INST] ")
 texts.append(f"{message.strip()} [/INST]")
 return "".join(texts)


class BackendType(Enum):
 UNKNOWN = 0
 TRANSFORMERS = 1
 GPTQ = 2
 LLAMA_CPP = 3

 @classmethod
 def get_type(cls, backend_name: str):
 backend_type = None
 backend_name_lower = backend_name.lower()
 if "transformers" in backend_name_lower:
 backend_type = BackendType.TRANSFORMERS
 elif "gptq" in backend_name_lower:
 backend_type = BackendType.GPTQ
 elif "cpp" in backend_name_lower:
 backend_type = BackendType.LLAMA_CPP
 else:
 raise Exception("Unknown backend: " + backend_name)
 # backend_type = BackendType.UNKNOWN
 return backend_type