app.py

import gradio as gr
from huggingface_hub import InferenceClient
import transformers
from transformers import AutoTokenizer,GenerationConfig, BitsAndBytesConfig
import torch
from peft import PeftModel
import spaces
import torch
import bitsandbytes, accelerate

print(transformers.__version__)  # Should be >= 4.26.0
print(bitsandbytes.__version__)  # Should be >= 0.37.0
print(accelerate.__version__)    # Should be >= 0.12.0

num_gpus = torch.cuda.device_count()
print(f"Number of available GPUs: {num_gpus}")

"""
For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
"""
#client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
base_model = "Neko-Institute-of-Science/LLaMA-65B-HF"
lora_weights = "./"
#lora_weights = LoraConfig(
#    auto_mapping=None,
#    base_model_name_or_path="Neko-Institute-of-Science/LLaMA-65B-HF",
#   bias=None,
#    fan_in_fan_out=False,
#    inference_mode=True,
#    init_lora_weights=True,
#    layers_pattern=None,
#    layers_to_transform=None,
#    lora_alpha=16,
#    lora_dropout=0.05,
#    modules_to_save=None,
#    peft_type="LORA",
#    revision=None,
#    target_modules=["q_proj","k_proj","v_proj","o_proj","gate_proj","up_proj","down_proj"],
#    task_type="CAUSAL_LM",
#)

cache_dir = "/data"

PROMPT_DICT = {
    "prompt_input": (
        "Below is an instruction that describes a task, paired with further context. "
        "Write a response that appropriately completes the request.\n\n"
        "Instruction:\n{instruction}\n\n Input:\n{input}\n\n Response:"
    ),
    "prompt_no_input": (
        "Below is an instruction that describes a task. "
        "Write a response that appropriately completes the request.\n\n"
        "Instruction:\n{instruction}\n\nResponse:"
    ),
}
model = None
tokenizer = None

quantization_config = BitsAndBytesConfig(
    load_in_8bit=True,  # Enable 8-bit quantization
    llm_int8_enable_fp32_cpu_offload=True  # Enable FP32 CPU offloading
)

def print_resources():
    # List details for each GPU
    for i in range(num_gpus):
        print(f"GPU {i}: {torch.cuda.get_device_name(i)}")
        print(f"  Total Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9:.2f} GB")
        print(f"  CUDA Capability: {torch.cuda.get_device_properties(i).major}.{torch.cuda.get_device_properties(i).minor}")
        print(f"  Allocated Memory: {torch.cuda.memory_allocated(i) / 1e9:.2f} GB")
        print(f"  Cached Memory: {torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
        print(f"  Free Memory: {torch.cuda.get_device_properties(i).total_memory / 1e9 - torch.cuda.memory_reserved(i) / 1e9:.2f} GB")
        
def generate_prompt(instruction, input=None):
    if input:
        return PROMPT_DICT["prompt_input"].format(instruction=instruction,input=input)
    else:
        return PROMPT_DICT["prompt_no_input"].format(instruction=instruction)

def generator(input_ids, generation_config, max_new_tokens):
    # Without streaming
    with torch.no_grad():
        generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=False,
                max_new_tokens=max_new_tokens,
            )
    return generation_output
    
def loadModel():
    global model, tokenizer
    if model is None:
        #from llama_rope_scaled_monkey_patch import replace_llama_rope_with_scaled_rope
        #replace_llama_rope_with_scaled_rope()
        model = transformers.AutoModelForCausalLM.from_pretrained(
                    base_model,
                    torch_dtype=torch.float16,
                    cache_dir=cache_dir,
                    device_map="auto",
                    #quantization_config=quantization_config,
                    max_memory={
                        0: "30GB",  # GPU 0 with 20GB memory
                        1: "45GB",  # GPU 0 with 20GB memory
                        2: "45GB",  # GPU 0 with 20GB memory
                        3: "45GB",  # GPU 0 with 20GB memory
                        #"cpu": "5GB"  # CPU with 100GB memory
                    },
                )
        print_resources()
        model = PeftModel.from_pretrained(
                    model,
                    lora_weights,
                    device_map="auto",
                    cache_dir='',
                    torch_dtype=torch.float16,
                    is_trainable=False,
                    max_memory={
                        0: "30GB",  # GPU 0 with 20GB memory
                        1: "45GB",  # GPU 0 with 20GB memory
                        2: "45GB",  # GPU 0 with 20GB memory
                        3: "45GB",  # GPU 0 with 20GB memory
                        #"cpu": "5GB"  # CPU with 100GB memory
                    },
                )
        tokenizer =  AutoTokenizer.from_pretrained(base_model,use_fast=False,cache_dir=cache_dir)
        tokenizer.pad_token = tokenizer.unk_token
    print_resources()
    return model, tokenizer

model, tokenizer = loadModel()

#@spaces.GPU(duration=120)
def respond(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
):
    ins_f = generate_prompt(message,None)
    inputs  =  tokenizer(ins_f, return_tensors="pt")
    print_resources()
    input_ids = inputs["input_ids"].cuda()
    max_new_tokens = 512
    generation_config = GenerationConfig(
            temperature=0.1,
            top_p=0.75,
            top_k=40,
            do_sample=True,
            num_beams=1,
            max_new_tokens = max_new_tokens
        )
    #generation_output = generator(input_ids, generation_config, max_new_tokens)
    with torch.no_grad():
        generation_output = model.generate(
                input_ids=input_ids,
                generation_config=generation_config,
                return_dict_in_generate=True,
                output_scores=False,
                max_new_tokens=max_new_tokens,
            )
    s = generation_output.sequences[0]
    output = tokenizer.decode(s)
    response = output.split("Response:")[1].strip()
    yield response
    
    #messages = [{"role": "system", "content": system_message}]

    #for val in history:
    #    if val[0]:
    #        messages.append({"role": "user", "content": val[0]})
    #    if val[1]:
    #        messages.append({"role": "assistant", "content": val[1]})

    # messages.append({"role": "user", "content": message})

    #response = ""

    #for message in client.chat_completion(
    #    messages,
    #    max_tokens=max_tokens,
    #    stream=True,
    #    temperature=temperature,
    #    top_p=top_p,
    #):
    #    token = message.choices[0].delta.content

    #    response += token
    #    yield response


"""
For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
"""
demo = gr.ChatInterface(
    respond,
    additional_inputs=[
        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
        gr.Slider(
            minimum=0.1,
            maximum=1.0,
            value=0.95,
            step=0.05,
            label="Top-p (nucleus sampling)",
        ),
    ],
)
if __name__ == "__main__":
    demo.launch()