Spaces:
Sleeping
Sleeping
| 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() | |