goavinash5's picture
Upload folder using huggingface_hub
e97665c
import os
import argparse
from typing import Iterator
import gradio as gr
from dotenv import load_dotenv
from distutils.util import strtobool
from llama2_wrapper import LLAMA2_WRAPPER
import logging
from prompts.utils import PromtsContainer
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--model_path", type=str, default="", help="model path")
parser.add_argument(
"--backend_type",
type=str,
default="",
help="Backend options: llama.cpp, gptq, transformers",
)
parser.add_argument(
"--load_in_8bit",
type=bool,
default=False,
help="Whether to use bitsandbytes 8 bit.",
)
parser.add_argument(
"--share",
type=bool,
default=False,
help="Whether to share public for gradio.",
)
args = parser.parse_args()
load_dotenv()
DEFAULT_SYSTEM_PROMPT = os.getenv("DEFAULT_SYSTEM_PROMPT", "")
MAX_MAX_NEW_TOKENS = int(os.getenv("MAX_MAX_NEW_TOKENS", 2048))
DEFAULT_MAX_NEW_TOKENS = int(os.getenv("DEFAULT_MAX_NEW_TOKENS", 1024))
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", 4000))
MODEL_PATH = os.getenv("MODEL_PATH")
assert MODEL_PATH is not None, f"MODEL_PATH is required, got: {MODEL_PATH}"
BACKEND_TYPE = os.getenv("BACKEND_TYPE")
assert BACKEND_TYPE is not None, f"BACKEND_TYPE is required, got: {BACKEND_TYPE}"
LOAD_IN_8BIT = bool(strtobool(os.getenv("LOAD_IN_8BIT", "True")))
if args.model_path != "":
MODEL_PATH = args.model_path
if args.backend_type != "":
BACKEND_TYPE = args.backend_type
if args.load_in_8bit:
LOAD_IN_8BIT = True
llama2_wrapper = LLAMA2_WRAPPER(
model_path=MODEL_PATH,
backend_type=BACKEND_TYPE,
max_tokens=MAX_INPUT_TOKEN_LENGTH,
load_in_8bit=LOAD_IN_8BIT,
# verbose=True,
)
DESCRIPTION = """
# llama2-webui
"""
DESCRIPTION2 = """
- Supporting models: [Llama-2-7b](https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML)/[13b](https://huggingface.co/llamaste/Llama-2-13b-chat-hf)/[70b](https://huggingface.co/llamaste/Llama-2-70b-chat-hf), [Llama-2-GPTQ](https://huggingface.co/TheBloke/Llama-2-7b-Chat-GPTQ), [Llama-2-GGML](https://huggingface.co/TheBloke/Llama-2-7B-Chat-GGML), [CodeLlama](https://huggingface.co/TheBloke/CodeLlama-7B-Instruct-GPTQ) ...
- Supporting model backends: [tranformers](https://github.com/huggingface/transformers), [bitsandbytes(8-bit inference)](https://github.com/TimDettmers/bitsandbytes), [AutoGPTQ(4-bit inference)](https://github.com/PanQiWei/AutoGPTQ), [llama.cpp](https://github.com/ggerganov/llama.cpp)
"""
def clear_and_save_textbox(message: str) -> tuple[str, str]:
return "", message
def save_textbox_for_prompt(message: str) -> str:
logging.info("start save_textbox_from_prompt")
message = convert_summary_to_prompt(message)
return message
def display_input(
message: str, history: list[tuple[str, str]]
) -> list[tuple[str, str]]:
history.append((message, ""))
return history
def delete_prev_fn(
history: list[tuple[str, str]]
) -> tuple[list[tuple[str, str]], str]:
try:
message, _ = history.pop()
except IndexError:
message = ""
return history, message or ""
def generate(
message: str,
history_with_input: list[tuple[str, str]],
system_prompt: str,
max_new_tokens: int,
temperature: float,
top_p: float,
top_k: int,
) -> Iterator[list[tuple[str, str]]]:
if max_new_tokens > MAX_MAX_NEW_TOKENS:
raise ValueError
try:
history = history_with_input[:-1]
generator = llama2_wrapper.run(
message,
history,
system_prompt,
max_new_tokens,
temperature,
top_p,
top_k,
)
try:
first_response = next(generator)
yield history + [(message, first_response)]
except StopIteration:
yield history + [(message, "")]
for response in generator:
yield history + [(message, response)]
except Exception as e:
logging.exception(e)
def check_input_token_length(
message: str, chat_history: list[tuple[str, str]], system_prompt: str
) -> None:
input_token_length = llama2_wrapper.get_input_token_length(
message, chat_history, system_prompt
)
if input_token_length > MAX_INPUT_TOKEN_LENGTH:
raise gr.Error(
f"The accumulated input is too long ({input_token_length} > {MAX_INPUT_TOKEN_LENGTH}). Clear your chat history and try again."
)
prompts_container = PromtsContainer()
prompts = prompts_container.get_prompts_tab_dict()
default_prompts_checkbox = False
default_advanced_checkbox = False
def convert_summary_to_prompt(summary):
return prompts_container.get_prompt_by_summary(summary)
def two_columns_list(tab_data, chatbot):
result = []
for i in range(int(len(tab_data) / 2) + 1):
row = gr.Row()
with row:
for j in range(2):
index = 2 * i + j
if index >= len(tab_data):
break
item = tab_data[index]
with gr.Group():
gr.HTML(
f'<p style="color: black; font-weight: bold;">{item["act"]}</p>'
)
prompt_text = gr.Button(
label="",
value=f"{item['summary']}",
size="sm",
elem_classes="text-left-aligned",
)
prompt_text.click(
fn=save_textbox_for_prompt,
inputs=prompt_text,
outputs=saved_input,
api_name=False,
queue=True,
).then(
fn=display_input,
inputs=[saved_input, chatbot],
outputs=chatbot,
api_name=False,
queue=True,
).then(
fn=check_input_token_length,
inputs=[saved_input, chatbot, system_prompt],
api_name=False,
queue=False,
).success(
fn=generate,
inputs=[
saved_input,
chatbot,
system_prompt,
max_new_tokens,
temperature,
top_p,
top_k,
],
outputs=chatbot,
api_name=False,
)
result.append(row)
return result
CSS = """
.contain { display: flex; flex-direction: column;}
#component-0 #component-1 #component-2 #component-4 #component-5 { height:71vh !important; }
#component-0 #component-1 #component-24 > div:nth-child(2) { height:80vh !important; overflow-y:auto }
.text-left-aligned {text-align: left !important; font-size: 16px;}
"""
with gr.Blocks(css=CSS) as demo:
with gr.Row(equal_height=True):
with gr.Column(scale=2):
gr.Markdown(DESCRIPTION)
with gr.Group():
chatbot = gr.Chatbot(label="Chatbot")
with gr.Row():
textbox = gr.Textbox(
container=False,
show_label=False,
placeholder="Type a message...",
scale=10,
)
submit_button = gr.Button(
"Submit", variant="primary", scale=1, min_width=0
)
with gr.Row():
retry_button = gr.Button("🔄 Retry", variant="secondary")
undo_button = gr.Button("↩️ Undo", variant="secondary")
clear_button = gr.Button("🗑️ Clear", variant="secondary")
saved_input = gr.State()
with gr.Row():
advanced_checkbox = gr.Checkbox(
label="Advanced",
value=default_prompts_checkbox,
container=False,
elem_classes="min_check",
)
prompts_checkbox = gr.Checkbox(
label="Prompts",
value=default_prompts_checkbox,
container=False,
elem_classes="min_check",
)
with gr.Column(visible=default_advanced_checkbox) as advanced_column:
system_prompt = gr.Textbox(
label="System prompt", value=DEFAULT_SYSTEM_PROMPT, lines=6
)
max_new_tokens = gr.Slider(
label="Max new tokens",
minimum=1,
maximum=MAX_MAX_NEW_TOKENS,
step=1,
value=DEFAULT_MAX_NEW_TOKENS,
)
temperature = gr.Slider(
label="Temperature",
minimum=0.1,
maximum=4.0,
step=0.1,
value=1.0,
)
top_p = gr.Slider(
label="Top-p (nucleus sampling)",
minimum=0.05,
maximum=1.0,
step=0.05,
value=0.95,
)
top_k = gr.Slider(
label="Top-k",
minimum=1,
maximum=1000,
step=1,
value=50,
)
with gr.Column(scale=1, visible=default_prompts_checkbox) as prompt_column:
gr.HTML(
'<p style="color: green; font-weight: bold;font-size: 16px;">\N{four leaf clover} prompts</p>'
)
for k, v in prompts.items():
with gr.Tab(k, scroll_to_output=True):
lst = two_columns_list(v, chatbot)
prompts_checkbox.change(
lambda x: gr.update(visible=x),
prompts_checkbox,
prompt_column,
queue=False,
)
advanced_checkbox.change(
lambda x: gr.update(visible=x),
advanced_checkbox,
advanced_column,
queue=False,
)
textbox.submit(
fn=clear_and_save_textbox,
inputs=textbox,
outputs=[textbox, saved_input],
api_name=False,
queue=False,
).then(
fn=display_input,
inputs=[saved_input, chatbot],
outputs=chatbot,
api_name=False,
queue=False,
).then(
fn=check_input_token_length,
inputs=[saved_input, chatbot, system_prompt],
api_name=False,
queue=False,
).success(
fn=generate,
inputs=[
saved_input,
chatbot,
system_prompt,
max_new_tokens,
temperature,
top_p,
top_k,
],
outputs=chatbot,
api_name=False,
)
button_event_preprocess = (
submit_button.click(
fn=clear_and_save_textbox,
inputs=textbox,
outputs=[textbox, saved_input],
api_name=False,
queue=False,
)
.then(
fn=display_input,
inputs=[saved_input, chatbot],
outputs=chatbot,
api_name=False,
queue=False,
)
.then(
fn=check_input_token_length,
inputs=[saved_input, chatbot, system_prompt],
api_name=False,
queue=False,
)
.success(
fn=generate,
inputs=[
saved_input,
chatbot,
system_prompt,
max_new_tokens,
temperature,
top_p,
top_k,
],
outputs=chatbot,
api_name=False,
)
)
retry_button.click(
fn=delete_prev_fn,
inputs=chatbot,
outputs=[chatbot, saved_input],
api_name=False,
queue=False,
).then(
fn=display_input,
inputs=[saved_input, chatbot],
outputs=chatbot,
api_name=False,
queue=False,
).then(
fn=generate,
inputs=[
saved_input,
chatbot,
system_prompt,
max_new_tokens,
temperature,
top_p,
top_k,
],
outputs=chatbot,
api_name=False,
)
undo_button.click(
fn=delete_prev_fn,
inputs=chatbot,
outputs=[chatbot, saved_input],
api_name=False,
queue=False,
).then(
fn=lambda x: x,
inputs=[saved_input],
outputs=textbox,
api_name=False,
queue=False,
)
clear_button.click(
fn=lambda: ([], ""),
outputs=[chatbot, saved_input],
queue=False,
api_name=False,
)
demo.queue(max_size=20).launch(share=args.share)
if __name__ == "__main__":
main()