add gradio file

attn_and_long_ctx_patches.py

@@ -0,0 +1,223 @@
+import torch
+from torch import nn
+from typing import Optional, Tuple, Union
+import transformers
+from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, rotate_half
+import math
+
+try:
+    from xformers import ops as xops
+except ImportError:
+    xops = None
+    print(
+        "Xformers is not installed correctly. If you want to use memory_efficient_attention use the following command to install Xformers\npip install xformers."
+    )
+
+
+STORE_KV_BEFORE_ROPE = False
+USE_MEM_EFF_ATTENTION = False
+ALPHA = 1.0
+AUTO_COEFF = 1.0
+SCALING_FACTOR = None
+
+
+def apply_rotary_pos_emb_single(q, cos, sin, position_ids):
+    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
+    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
+    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]
+    q_embed = (q * cos) + (rotate_half(q) * sin)
+    return q_embed
+
+
+def xformers_forward(
+    self,
+    hidden_states: torch.Tensor,
+    attention_mask: Optional[torch.Tensor] = None,
+    position_ids: Optional[torch.LongTensor] = None,
+    past_key_value: Optional[Tuple[torch.Tensor]] = None,
+    output_attentions: bool = False,
+    use_cache: bool = False,
+) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
+    bsz, q_len, _ = hidden_states.size()
+
+    query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+    value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
+
+    kv_seq_len = key_states.shape[-2]
+    if past_key_value is not None:
+        kv_seq_len += past_key_value[0].shape[-2]
+
+    if STORE_KV_BEFORE_ROPE is False:
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
+        # [bsz, nh, t, hd]
+
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+
+        past_key_value = (key_states, value_states) if use_cache else None
+    else:
+        if past_key_value is not None:
+            # reuse k, v, self_attention
+            key_states = torch.cat([past_key_value[0], key_states], dim=2)
+            value_states = torch.cat([past_key_value[1], value_states], dim=2)
+        past_key_value = (key_states, value_states) if use_cache else None
+
+        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)
+
+        query_states = apply_rotary_pos_emb_single(query_states, cos, sin, position_ids)
+        position_ids = torch.arange(kv_seq_len, dtype=torch.long, device=cos.device)
+        position_ids = position_ids.unsqueeze(0).view(-1, kv_seq_len)
+        key_states = apply_rotary_pos_emb_single(key_states, cos, sin, position_ids)
+
+    if xops is not None and USE_MEM_EFF_ATTENTION:
+        attn_weights = None
+        query_states = query_states.transpose(1, 2)
+        key_states = key_states.transpose(1, 2)
+        value_states = value_states.transpose(1, 2)
+        attn_bias = None if (query_states.size(1)==1 and key_states.size(1)>1) else xops.LowerTriangularMask()
+        attn_output = xops.memory_efficient_attention(
+            query_states, key_states, value_states, attn_bias=attn_bias, p=0)
+    else:
+        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
+
+        if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):
+            raise ValueError(
+                f"Attention weights should be of size {(bsz * self.num_heads, q_len, kv_seq_len)}, but is"
+                f" {attn_weights.size()}"
+            )
+
+        if attention_mask is not None:
+            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):
+                raise ValueError(
+                    f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"
+                )
+            attn_weights = attn_weights + attention_mask
+            attn_weights = torch.max(
+                attn_weights, torch.tensor(torch.finfo(attn_weights.dtype).min, device=attn_weights.device)
+            )
+
+        # upcast attention to fp32
+        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
+        attn_output = torch.matmul(attn_weights, value_states)
+
+        if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):
+            raise ValueError(
+                f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"
+                f" {attn_output.size()}"
+            )
+
+        attn_output = attn_output.transpose(1, 2)
+    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
+
+    attn_output = self.o_proj(attn_output)
+
+    if not output_attentions:
+        attn_weights = None
+
+    return attn_output, attn_weights, past_key_value
+
+
+old_init = transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__
+
+
+def _set_cos_sin_cache(self, seq_len, device, dtype):
+    self.max_seq_len_cached = seq_len
+    t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)
+    t = t / self.scaling_factor
+
+    freqs = torch.einsum("i,j->ij", t, self.ntk_inv_freq.to(device))
+    # Different from paper, but it uses a different permutation in order to obtain the same calculation
+    emb = torch.cat((freqs, freqs), dim=-1)
+    self.register_buffer("cos_cached", emb.cos()[None, None, :, :], persistent=False)
+    self.register_buffer("sin_cached", emb.sin()[None, None, :, :], persistent=False)
+
+
+def adaptive_ntk_init(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=None):
+    self.alpha = ALPHA
+    if SCALING_FACTOR is None:
+        self.scaling_factor = scaling_factor or 1.0
+    else:
+        self.scaling_factor = SCALING_FACTOR
+    if isinstance(ALPHA,(float,int)):
+        base = base * ALPHA ** (dim / (dim-2))
+        self.base = base
+    elif ALPHA=='auto':
+        self.base = base
+    else:
+        raise ValueError(ALPHA)
+    old_init(self, dim, max_position_embeddings, base, device)
+    self.ntk_inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(device) / dim))
+
+    self._set_cos_sin_cache = _set_cos_sin_cache
+    self._set_cos_sin_cache(
+        self, seq_len=max_position_embeddings, device=self.ntk_inv_freq.device, dtype=torch.get_default_dtype()
+    )
+
+
+def adaptive_ntk_forward(self, x, seq_len=None):
+    if seq_len > self.max_seq_len_cached:
+        if isinstance(self.alpha,(float,int)):
+            self._set_cos_sin_cache(self, seq_len=seq_len, device=x.device, dtype=x.dtype)
+        elif self.alpha=='auto':
+            t = torch.arange(seq_len, device=x.device, dtype=torch.float32)
+            t = t / self.scaling_factor
+            dim = self.dim
+            alpha = (seq_len / (self.max_position_embeddings/2) - 1) * AUTO_COEFF
+            base = self.base * alpha ** (dim / (dim-2))
+            ntk_inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float().to(x.device) / dim ))
+
+            freqs = torch.einsum("i,j->ij", t, ntk_inv_freq)
+            emb = torch.cat((freqs, freqs), dim=-1).to(x.device)
+            cos_cached = emb.cos()[None, None, :, :]
+            sin_cached = emb.sin()[None, None, :, :]
+            return (
+                cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+                sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype)
+            )
+    return (
+        self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),
+        self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype)
+    )
+
+
+def apply_attention_patch(
+        use_memory_efficient_attention=False,
+        store_kv_before_rope=False
+        ):
+    global USE_MEM_EFF_ATTENTION, STORE_KV_BEFORE_ROPE
+    if use_memory_efficient_attention is True and xops is not None:
+        USE_MEM_EFF_ATTENTION = use_memory_efficient_attention
+    print("USE_MEM_EFF_ATTENTION: ",USE_MEM_EFF_ATTENTION)
+    STORE_KV_BEFORE_ROPE = store_kv_before_rope
+    print("STORE_KV_BEFORE_ROPE:", STORE_KV_BEFORE_ROPE)
+    transformers.models.llama.modeling_llama.LlamaAttention.forward = xformers_forward
+
+
+def apply_ntk_scaling_patch(alpha: Union[float,str], scaling_factor: Optional[float] = None):
+    global ALPHA
+    global SCALING_FACTOR
+    ALPHA = alpha
+    SCALING_FACTOR = scaling_factor
+    try:
+        ALPHA = float(ALPHA)
+    except ValueError:
+        if ALPHA!="auto":
+            raise ValueError(f"Alpha can only be a float or 'auto', but given {ALPHA}")
+    print(f"Apply NTK scaling with ALPHA={ALPHA}")
+    if scaling_factor is None:
+        print(f"The value of scaling factor will be read from model config file, or set to 1.")
+    else:
+        print(f"Warning: scaling factor is set to {SCALING_FACTOR}. \
+              If you set the value by hand, do not forget to update \
+              max_position_embeddings in the model config file.")
+
+    transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.__init__ = adaptive_ntk_init
+    if hasattr(transformers.models.llama.modeling_llama,'LlamaLinearScalingRotaryEmbedding'):
+        transformers.models.llama.modeling_llama.LlamaLinearScalingRotaryEmbedding.__init__ = adaptive_ntk_init
+    transformers.models.llama.modeling_llama.LlamaRotaryEmbedding.forward = adaptive_ntk_forward

gradio_demo.py

@@ -0,0 +1,626 @@
+import torch
+from transformers import (
+    LlamaForCausalLM,
+    LlamaTokenizer,
+    StoppingCriteria,
+    BitsAndBytesConfig
+)
+import gradio as gr
+import argparse
+import os
+from queue import Queue
+from threading import Thread
+import traceback
+import gc
+import json
+import requests
+from typing import Iterable, List
+import subprocess
+import re
+
+DEFAULT_SYSTEM_PROMPT = """You are a helpful, respectful and honest assistant.Help as much as you can."""
+
+TEMPLATE_WITH_SYSTEM_PROMPT = (
+    "[INST] <<SYS>>\n"
+    "{system_prompt}\n"
+    "<</SYS>>\n\n"
+    "{instruction} [/INST]"
+)
+
+TEMPLATE_WITHOUT_SYSTEM_PROMPT = "[INST] {instruction} [/INST]"
+
+# Parse command-line arguments
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    '--base_model',
+    default=None,
+    type=str,
+    required=True,
+    help='Base model path')
+parser.add_argument('--lora_model', default=None, type=str,
+                    help="If None, perform inference on the base model")
+parser.add_argument(
+    '--tokenizer_path',
+    default=None,
+    type=str,
+    help='If None, lora model path or base model path will be used')
+parser.add_argument(
+    '--gpus',
+    default="0",
+    type=str,
+    help='If None, cuda:0 will be used. Inference using multi-cards: --gpus=0,1,... ')
+parser.add_argument('--share', default=True, help='Share gradio domain name')
+parser.add_argument('--port', default=19324, type=int, help='Port of gradio demo')
+parser.add_argument(
+    '--max_memory',
+    default=1024,
+    type=int,
+    help='Maximum number of input tokens (including system prompt) to keep. If exceeded, earlier history will be discarded.')
+parser.add_argument(
+    '--load_in_8bit',
+    action='store_true',
+     default=False,
+    help='Use 8 bit quantized model')
+parser.add_argument(
+    '--load_in_4bit',
+    action='store_true',
+    default=False,
+    help='Use 4 bit quantized model')
+parser.add_argument(
+    '--only_cpu',
+    action='store_true',
+    help='Only use CPU for inference')
+parser.add_argument(
+    '--alpha',
+    type=str,
+    default="1.0",
+    help="The scaling factor of NTK method, can be a float or 'auto'. ")
+parser.add_argument(
+    "--use_vllm",
+    action='store_true',
+    help="Use vLLM as back-end LLM service.")
+parser.add_argument(
+    "--post_host",
+    type=str,
+    default="0.0.0.0",
+    help="Host of vLLM service.")
+parser.add_argument(
+    "--post_port",
+    type=int,
+    default=7777,
+    help="Port of vLLM service.")
+args = parser.parse_args()
+
+ENABLE_CFG_SAMPLING = True
+try:
+    from transformers.generation import UnbatchedClassifierFreeGuidanceLogitsProcessor
+except ImportError:
+    ENABLE_CFG_SAMPLING = False
+    print("Install the latest transformers (commit equal or later than d533465) to enable CFG sampling.")
+if args.use_vllm is True:
+    print("CFG sampling is disabled when using vLLM.")
+    ENABLE_CFG_SAMPLING = False
+
+if args.only_cpu is True:
+    args.gpus = ""
+    if args.load_in_8bit or args.load_in_4bit:
+        raise ValueError("Quantization is unavailable on CPU.")
+if args.load_in_8bit and args.load_in_4bit:
+    raise ValueError("Only one quantization method can be chosen for inference. Please check your arguments")
+import sys
+parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
+sys.path.append(parent_dir)
+from attn_and_long_ctx_patches import apply_attention_patch, apply_ntk_scaling_patch
+if not args.only_cpu:
+    apply_attention_patch(use_memory_efficient_attention=True)
+apply_ntk_scaling_patch(args.alpha)
+
+# Set CUDA devices if available
+os.environ["CUDA_VISIBLE_DEVICES"] = args.gpus
+
+
+# Peft library can only import after setting CUDA devices
+from peft import PeftModel
+
+
+# Set up the required components: model and tokenizer
+
+def setup():
+    global tokenizer, model, device, share, port, max_memory
+    if args.use_vllm:
+        # global share, port, max_memory
+        max_memory = args.max_memory
+        port = args.port
+        share = args.share
+
+        if args.lora_model is not None:
+            raise ValueError("vLLM currently does not support LoRA, please merge the LoRA weights to the base model.")
+        if args.load_in_8bit or args.load_in_4bit:
+            raise ValueError("vLLM currently does not support quantization, please use fp16 (default) or unuse --use_vllm.")
+        if args.only_cpu:
+            raise ValueError("vLLM requires GPUs with compute capability not less than 7.0. If you want to run only on CPU, please unuse --use_vllm.")
+
+        if args.tokenizer_path is None:
+            args.tokenizer_path = args.base_model
+        tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
+
+        print("Start launch vllm server.")
+        cmd = f"python -m vllm.entrypoints.api_server \
+            --model={args.base_model} \
+            --tokenizer={args.tokenizer_path} \
+            --tokenizer-mode=slow \
+            --tensor-parallel-size={len(args.gpus.split(','))} \
+            --host {args.post_host} \
+            --port {args.post_port} \
+            &"
+        subprocess.check_call(cmd, shell=True)
+    else:
+        max_memory = args.max_memory
+        port = args.port
+        share = args.share
+        load_type = torch.float16
+        if torch.cuda.is_available():
+            device = torch.device(0)
+        else:
+            device = torch.device('cpu')
+        if args.tokenizer_path is None:
+            args.tokenizer_path = args.base_model
+            # if args.lora_model is None:
+            #     args.tokenizer_path = args.base_model
+        tokenizer = LlamaTokenizer.from_pretrained(args.tokenizer_path, legacy=True)
+        tokenizer.pad_token_id = 0
+        # tokenizer.pad_token = "<>"
+        base_model = LlamaForCausalLM.from_pretrained(
+            args.base_model,
+            torch_dtype=load_type,
+            low_cpu_mem_usage=True,
+            device_map='auto',
+            quantization_config=BitsAndBytesConfig(
+                load_in_4bit=args.load_in_4bit,
+                load_in_8bit=args.load_in_8bit,
+                bnb_4bit_compute_dtype=load_type,
+                # load_in_8bit_fp32_cpu_offload=True
+            )
+        )
+
+        model_vocab_size = base_model.get_input_embeddings().weight.size(0)
+        tokenizer_vocab_size = len(tokenizer)
+        print(f"Vocab of the base model: {model_vocab_size}")
+        print(f"Vocab of the tokenizer: {tokenizer_vocab_size}")
+        if model_vocab_size != tokenizer_vocab_size:
+            print("Resize model embeddings to fit tokenizer")
+            base_model.resize_token_embeddings(tokenizer_vocab_size)
+        if args.lora_model is not None:
+            print("loading peft model")
+            model = PeftModel.from_pretrained(
+                base_model,
+                args.lora_model,
+                torch_dtype=load_type,
+                device_map='auto',
+            ).half()
+        else:
+            model = base_model
+
+        if device == torch.device('cpu'):
+            model.float()
+
+        model.eval()
+
+
+# Reset the user input
+def reset_user_input():
+    return gr.update(value='')
+
+
+# Reset the state
+def reset_state():
+    return []
+
+
+def generate_prompt(instruction, response="", with_system_prompt=True, system_prompt=DEFAULT_SYSTEM_PROMPT):
+    if with_system_prompt is True:
+        prompt = TEMPLATE_WITH_SYSTEM_PROMPT.format_map({'instruction': instruction,'system_prompt': system_prompt})
+    else:
+        prompt = TEMPLATE_WITHOUT_SYSTEM_PROMPT.format_map({'instruction': instruction})
+    if len(response)>0:
+        prompt += " " + response
+    return prompt
+
+
+# User interaction function for chat
+def user(user_message, history):
+    return gr.update(value="", interactive=False), history + \
+        [[user_message, None]]
+
+
+class Stream(StoppingCriteria):
+    def __init__(self, callback_func=None):
+        self.callback_func = callback_func
+
+    def __call__(self, input_ids, scores) -> bool:
+        if self.callback_func is not None:
+            self.callback_func(input_ids[0])
+        return False
+
+
+class Iteratorize:
+    """
+    Transforms a function that takes a callback
+    into a lazy iterator (generator).
+
+    Adapted from: https://stackoverflow.com/a/9969000
+    """
+    def __init__(self, func, kwargs=None, callback=None):
+        self.mfunc = func
+        self.c_callback = callback
+        self.q = Queue()
+        self.sentinel = object()
+        self.kwargs = kwargs or {}
+        self.stop_now = False
+
+        def _callback(val):
+            if self.stop_now:
+                raise ValueError
+            self.q.put(val)
+
+        def gentask():
+            try:
+                ret = self.mfunc(callback=_callback, **self.kwargs)
+            except ValueError:
+                pass
+            except Exception:
+                traceback.print_exc()
+
+            clear_torch_cache()
+            self.q.put(self.sentinel)
+            if self.c_callback:
+                self.c_callback(ret)
+
+        self.thread = Thread(target=gentask)
+        self.thread.start()
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        obj = self.q.get(True, None)
+        if obj is self.sentinel:
+            raise StopIteration
+        else:
+            return obj
+
+    def __del__(self):
+        clear_torch_cache()
+
+    def __enter__(self):
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        self.stop_now = True
+        clear_torch_cache()
+
+
+def clear_torch_cache():
+    gc.collect()
+    if torch.cuda.device_count() > 0:
+        torch.cuda.empty_cache()
+
+
+def post_http_request(prompt: str,
+                      api_url: str,
+                      n: int = 1,
+                      top_p: float = 0.9,
+                      top_k: int = 40,
+                      temperature: float = 0.2,
+                      max_tokens: int = 1024,
+                      presence_penalty: float = 1.0,
+                      use_beam_search: bool = False,
+                      stream: bool = False) -> requests.Response:
+    headers = {"User-Agent": "Test Client"}
+    pload = {
+        "prompt": prompt,
+        "n": n,
+        "top_p": 1 if use_beam_search else top_p,
+        "top_k": -1 if use_beam_search else top_k,
+        "temperature": 0 if use_beam_search else temperature,
+        "max_tokens": max_tokens,
+        "use_beam_search": use_beam_search,
+        "best_of": 5 if use_beam_search else n,
+        "presence_penalty": presence_penalty,
+        "stream": stream,
+    }
+    print(pload)
+
+    response = requests.post(api_url, headers=headers, json=pload, stream=True)
+    return response
+
+
+def get_streaming_response(response: requests.Response) -> Iterable[List[str]]:
+    for chunk in response.iter_lines(chunk_size=8192,
+                                     decode_unicode=False,
+                                     delimiter=b"\0"):
+        if chunk:
+            data = json.loads(chunk.decode("utf-8"))
+            output = data["text"]
+            yield output
+
+
+# Perform prediction based on the user input and history
+@torch.no_grad()
+def predict(
+    history,
+    system_prompt,
+    negative_prompt,
+    max_new_tokens=1024,
+    top_p=0.89,
+    temperature=0.85,
+    top_k=40,
+    do_sample=True,
+    repetition_penalty=1.2,
+    guidance_scale=1.0,
+    presence_penalty=0.0,
+):
+    if len(system_prompt) == 0:
+        system_prompt = DEFAULT_SYSTEM_PROMPT
+    while True:
+        print("len(history):", len(history))
+        print("history: ", history)
+        history[-1][1] = ""
+        if len(history) == 1:
+            input = history[0][0]
+            prompt = generate_prompt(input,response="", with_system_prompt=True, system_prompt=system_prompt)
+            print(f"prompt:{prompt}")
+        else:
+            input = history[0][0]
+            response = history[0][1]
+            prompt = generate_prompt(input, response=response, with_system_prompt=True, system_prompt=system_prompt)+'</s>'
+            for hist in history[1:-1]:
+                input = hist[0]
+                response = hist[1]
+                prompt = prompt + '<s>'+generate_prompt(input, response=response, with_system_prompt=False)+'</s>'
+            input = history[-1][0]
+            check_text = input.replace("<br>","").replace(" ","").replace("\n","")
+            if len(check_text) == 0:
+                input = ""
+            prompt = prompt + '<s>'+generate_prompt(input, response="", with_system_prompt=False)
+            print(f"prompt1:{prompt}")
+        input_length = len(tokenizer.encode(prompt, add_special_tokens=True))
+        print(f"Input length: {input_length}")
+        if input_length > max_memory and len(history) > 1:
+            print(f"The input length ({input_length}) exceeds the max memory ({max_memory}). The earlier history will be discarded.")
+            history = history[1:]
+            print("history: ", history)
+        else:
+            break
+
+    if args.use_vllm:
+        generate_params = {
+            'max_tokens': max_new_tokens,
+            'top_p': top_p,
+            'temperature': temperature,
+            'top_k': top_k,
+            "use_beam_search": not do_sample,
+            'presence_penalty': presence_penalty,
+        }
+
+        api_url = f"http://{args.post_host}:{args.post_port}/generate"
+
+
+        response = post_http_request(prompt, api_url, **generate_params, stream=True)
+
+        for h in get_streaming_response(response):
+            for line in h:
+                line = line.replace(prompt, '')
+                history[-1][1] = line
+                yield history
+
+    else:
+        negative_text = None
+        if len(negative_prompt) != 0:
+            negative_text = re.sub(r"<<SYS>>\n(.*)\n<</SYS>>", f"<<SYS>>\n{negative_prompt}\n<</SYS>>", prompt)
+        inputs = tokenizer(prompt, return_tensors="pt")
+        input_ids = inputs["input_ids"].to(device)
+        if negative_text is None:
+            negative_prompt_ids = None
+            negative_prompt_attention_mask = None
+        else:
+            negative_inputs = tokenizer(negative_text,return_tensors="pt")
+            negative_prompt_ids = negative_inputs["input_ids"].to(device)
+            negative_prompt_attention_mask = negative_inputs["attention_mask"].to(device)
+        generate_params = {
+            'input_ids': input_ids,
+            'max_new_tokens': max_new_tokens,
+            'top_p': top_p,
+            'temperature': temperature,
+            'top_k': top_k,
+            'do_sample': do_sample,
+            'repetition_penalty': repetition_penalty,
+        }
+        if ENABLE_CFG_SAMPLING is True:
+            generate_params['guidance_scale'] = guidance_scale
+            generate_params['negative_prompt_ids'] = negative_prompt_ids
+            generate_params['negative_prompt_attention_mask'] = negative_prompt_attention_mask
+
+        def generate_with_callback(callback=None, **kwargs):
+            if 'stopping_criteria' in kwargs:
+                kwargs['stopping_criteria'].append(Stream(callback_func=callback))
+            else:
+                kwargs['stopping_criteria'] = [Stream(callback_func=callback)]
+            clear_torch_cache()
+            with torch.no_grad():
+                model.generate(**kwargs)
+
+        def generate_with_streaming(**kwargs):
+            return Iteratorize(generate_with_callback, kwargs, callback=None)
+
+        with generate_with_streaming(**generate_params) as generator:
+            for output in generator:
+                next_token_ids = output[len(input_ids[0]):]
+                if next_token_ids[0] in [tokenizer.eos_token_id,0]:
+                    break
+                new_tokens = tokenizer.decode(
+                    next_token_ids, skip_special_tokens=True)
+                if isinstance(tokenizer, LlamaTokenizer) and len(next_token_ids) > 0:
+                    if tokenizer.convert_ids_to_tokens(int(next_token_ids[0])).startswith('▁'):
+                        new_tokens = ' ' + new_tokens
+
+                history[-1][1] = new_tokens
+                yield history
+                if len(next_token_ids) >= max_new_tokens:
+                    break
+
+
+# Call the setup function to initialize the components
+setup()
+
+
+# Create the Gradio interface
+with gr.Blocks(
+    theme=gr.themes.Soft(),
+    css=".disclaimer {font-variant-caps: all-small-caps;}") as demo:
+    github_banner_path = 'https://raw.githubusercontent.com/moseshu/llama2-chat/main/llama2.jpg'
+    gr.HTML(f'<p align="center"><a href="https://huggingface.co/Moses25/Llama2-Moses-7b-chat"><img src={github_banner_path} width="100" height="40"/>Llama2-Moses-7b</a></p>')
+    chatbot = gr.Chatbot().style(height=300)
+    with gr.Row():
+        with gr.Column(scale=4):
+            with gr.Column(scale=3):
+                system_prompt_input = gr.Textbox(
+                    show_label=True,
+                    label="system prompt(仅在对话开始前或清空历史后修改有效，对话过程中修改无效)",
+                    placeholder=DEFAULT_SYSTEM_PROMPT,
+                    lines=1).style(
+                    container=True)
+                negative_prompt_input = gr.Textbox(
+                    show_label=True,
+                    label="反向提示语（仅在对话开始前或清空历史后修改有效，对话过程中修改无效）",
+                    placeholder="option",
+                    lines=1,
+                    visible=ENABLE_CFG_SAMPLING).style(
+                    container=True)
+            with gr.Column(scale=10):
+                user_input = gr.Textbox(
+                    show_label=True,
+                    label="ChatBox",
+                    placeholder="Shift + Enter发送消息...",
+                    lines=10).style(
+                    container=True)
+            with gr.Column(min_width=24, scale=1):
+                with gr.Row():
+                    stop = gr.Button("Stop",variant='stop')
+                    submitBtn = gr.Button("Submit", variant="primary")
+        with gr.Column(scale=1):
+            emptyBtn = gr.Button("Clear History")
+            max_new_token = gr.Slider(
+                0,
+                4096,
+                value=1024,
+                step=1.0,
+                label="Maximum New Token Length",
+                interactive=True)
+            top_p = gr.Slider(0, 1, value=0.9, step=0.01,
+                              label="Top P", interactive=True)
+            temperature = gr.Slider(
+                0,
+                1,
+                value=0.7,
+                step=0.01,
+                label="Temperature",
+                interactive=True)
+            top_k = gr.Slider(1, 40, value=40, step=1,
+                              label="Top K", interactive=True)
+            do_sample = gr.Checkbox(
+                value=True,
+                label="Do Sample",
+                info="use random sample strategy",
+                interactive=True)
+            repetition_penalty = gr.Slider(
+                1.0,
+                3.0,
+                value=1.1,
+                step=0.1,
+                label="Repetition Penalty",
+                interactive=True,
+                visible=False if args.use_vllm else True)
+            guidance_scale = gr.Slider(
+                1.0,
+                3.0,
+                value=1.0,
+                step=0.1,
+                label="Guidance Scale",
+                interactive=True,
+                visible=ENABLE_CFG_SAMPLING)
+            presence_penalty = gr.Slider(
+                -2.0,
+                2.0,
+                value=1.0,
+                step=0.1,
+                label="Presence Penalty",
+                interactive=True,
+                visible=True if args.use_vllm else False)
+
+    
+    params = [user_input, chatbot]
+    predict_params = [
+        chatbot,
+        system_prompt_input,
+        negative_prompt_input,
+        max_new_token,
+        top_p,
+        temperature,
+        top_k,
+        do_sample,
+        repetition_penalty,
+        guidance_scale,
+        presence_penalty]
+    with gr.Row():
+        gr.Markdown(
+            "免责声明：该模型可能会产生与事实不符的输出，不应依赖该模型来产生与事实相符的信息。模型在各种公共数据集以及得物一些商品信息进行训练。尽管做了大量的数据清洗，但是模型的输出结果还可能存在一些问题",
+            elem_classes=["disclaimer"],
+        )
+    submit_click_event = submitBtn.click(
+        user,
+        params,
+        params,
+        queue=False).then(
+        predict,
+        predict_params,
+        chatbot).then(
+            lambda: gr.update(
+                interactive=True),
+        None,
+        [user_input],
+        queue=True)
+
+    
+    submit_event = user_input.submit(
+        user,
+        params,
+        params,
+        queue=False).then(
+        predict,
+        predict_params,
+        chatbot).then(
+            lambda: gr.update(
+                interactive=True),
+        None,
+        [user_input],
+        queue=True)
+
+    submitBtn.click(reset_user_input, [], [user_input])
+    
+    stop.click(
+        fn=None,
+        inputs=None,
+        outputs=None,
+        cancels=[submit_event, submit_click_event],
+        queue=False,
+    )
+    emptyBtn.click(reset_state, outputs=[chatbot], show_progress=True)
+
+
+# Launch the Gradio interface
+demo.queue().launch(
+    share=share,
+    inbrowser=True,
+    server_name='0.0.0.0',
+    server_port=port)