Update app.py

app.py

@@ -1,259 +1,263 @@
-from pathlib import Path
-from shutil import rmtree
-from typing import Union, List, Dict, Tuple, Optional
-from tqdm import tqdm
-
-import requests
-import gradio as gr
-from llama_cpp import Llama
-
-
-# ================== ANNOTATIONS ========================
-
-CHAT_HISTORY = List[Tuple[Optional[str], Optional[str]]]
-MODEL_DICT = Dict[str, Llama]
-
-
-# ================== FUNCS =============================
-
-def download_file(file_url: str, file_path: Union[str, Path]) -> None:
-    response = requests.get(file_url, stream=True)
-    if response.status_code != 200:
-        raise Exception(f'Файл недоступен для скачивания по ссылке: {file_url}')
-    total_size = int(response.headers.get('content-length', 0))
-    progress_tqdm = tqdm(desc='Loading GGUF file', total=total_size, unit='iB', unit_scale=True)
-    progress_gradio = gr.Progress()
-    completed_size = 0
-    with open(file_path, 'wb') as file:
-        for data in response.iter_content(chunk_size=4096):
-            size = file.write(data)
-            progress_tqdm.update(size)
-            completed_size += size
-            desc = f'Loading GGUF file, {completed_size/1024**3:.3f}/{total_size/1024**3:.3f} GB'
-            progress_gradio(completed_size/total_size, desc=desc)
-
-
-def download_gguf_and_init_model(gguf_url: str, model_dict: MODEL_DICT) -> Tuple[MODEL_DICT, bool, str]:
-    log = ''
-    if not gguf_url.endswith('.gguf'):
-        log += f'The link must be a direct link to the GGUF file\n'
-        return model_dict, log
-
-    gguf_filename = gguf_url.rsplit('/')[-1]
-    model_path = MODELS_PATH / gguf_filename
-    progress = gr.Progress()
-
-    if not model_path.is_file():
-        progress(0.3, desc='Шаг 1/2: Loading GGUF model file')
-        try:
-            download_file(gguf_url, model_path)
-            log += f'Model file {gguf_filename} successfully loaded\n'
-        except Exception as ex:
-            log += f'Error loading model from link {gguf_url}, error code:\n{ex}\n'
-            curr_model = model_dict.get('model')
-            if curr_model is None:
-                log += f'Model is missing from dictionary "model_dict"\n'
-                return model_dict, load_log
-            curr_model_filename = Path(curr_model.model_path).name
-            log += f'Current initialized model: {curr_model_filename}\n'
-            return model_dict, log
-    else:
-        log += f'Model file {gguf_filename} loaded, initializing model...\n'
-
-    progress(0.7, desc='Шаг 2/2: Model initialization')
-    model = Llama(model_path=str(model_path), n_gpu_layers=-1, verbose=True)
-    model_dict = {'model': model}
-    support_system_role = 'System role not supported' not in model.metadata['tokenizer.chat_template']
-    log += f'Model {gguf_filename} initialized\n'
-    return model_dict, support_system_role, log
-
-
-def user_message_to_chatbot(user_message: str, chatbot: CHAT_HISTORY) -> Tuple[str, CHAT_HISTORY]:
-    if user_message:
-        chatbot.append((user_message, None))
-    return '', chatbot
-
-
-def bot_response_to_chatbot(
-        chatbot: CHAT_HISTORY,
-        model_dict: MODEL_DICT,
-        system_prompt: str,
-        support_system_role: bool,
-        history_len: int,
-        do_sample: bool,
-        *generate_args,
-        ):
-
-    model = model_dict.get('model')
-    user_message = chatbot[-1][0]
-    messages = []
-
-    gen_kwargs = dict(zip(GENERATE_KWARGS.keys(), generate_args))
-    gen_kwargs['top_k'] = int(gen_kwargs['top_k'])
-
-    if not do_sample:
-        gen_kwargs['top_p'] = 0.0
-        gen_kwargs['top_k'] = 1
-        gen_kwargs['repeat_penalty'] = 1.0
-
-    if support_system_role and system_prompt:
-        messages.append({'role': 'system', 'content': system_prompt})
-
-    if history_len != 0:
-        for user_msg, bot_msg in chatbot[:-1][-history_len:]:
-            print(user_msg, bot_msg)
-            messages.append({'role': 'user', 'content': user_msg})
-            messages.append({'role': 'assistant', 'content': bot_msg})
-
-    messages.append({'role': 'user', 'content': user_message})
-    stream_response = model.create_chat_completion(
-        messages=messages,
-        stream=True,
-        **gen_kwargs,
-        )
-
-    chatbot[-1][1] = ''
-    for chunk in stream_response:
-        token = chunk['choices'][0]['delta'].get('content')
-        if token is not None:
-            chatbot[-1][1] += token
-            yield chatbot
-
-
-def get_system_prompt_component(interactive: bool) -> gr.Textbox:
-    value = '' if interactive else 'System prompt is not supported by this model'
-    return gr.Textbox(value=value, label='System prompt', interactive=interactive)
-
-
-def get_generate_args(do_sample: bool) -> List[gr.component]:
-    visible = do_sample
-    generate_args = [
-        gr.Slider(label='temperature', value=GENERATE_KWARGS['temperature'], minimum=0.1, maximum=3, step=0.1, visible=visible),
-        gr.Slider(label='top_p', value=GENERATE_KWARGS['top_p'], minimum=0.1, maximum=1, step=0.1, visible=visible),
-        gr.Slider(label='top_k', value=GENERATE_KWARGS['top_k'], minimum=1, maximum=50, step=5, visible=visible),
-        gr.Slider(label='repeat_penalty', value=GENERATE_KWARGS['repeat_penalty'], minimum=1, maximum=5, step=0.1, visible=visible),
-    ]
-    return generate_args
-
-
-# ================== VARIABLES =============================
-
-MODELS_PATH = Path('models')
-MODELS_PATH.mkdir(exist_ok=True)
-DEFAULT_GGUF_URL = 'https://huggingface.co/bartowski/gemma-2-2b-it-GGUF/resolve/main/gemma-2-2b-it-Q8_0.gguf'
-
-start_model_dict, start_support_system_role, start_load_log = download_gguf_and_init_model(
-    gguf_url=DEFAULT_GGUF_URL, model_dict={},
-    )
-
-GENERATE_KWARGS = dict(
-    temperature=0.2,
-    top_p=0.95,
-    top_k=40,
-    repeat_penalty=1.0,
-    )
-
-theme = gr.themes.Base(primary_hue='green', secondary_hue='yellow', neutral_hue='zinc').set(
-    loader_color='rgb(0, 255, 0)',
-    slider_color='rgb(0, 200, 0)',
-    body_text_color_dark='rgb(0, 200, 0)',
-    button_secondary_background_fill_dark='green',
-)
-css = '''.gradio-container {width: 60% !important}'''
-
-
-# ================== INTERFACE =============================
-
-with gr.Blocks(theme=theme, css=css) as interface:
-    model_dict = gr.State(start_model_dict)
-    support_system_role = gr.State(start_support_system_role)
-    
-    # ================= CHAT BOT PAGE ======================
-    with gr.Tab('Chat bot'):
-        with gr.Row():
-            with gr.Column(scale=3):
-                chatbot = gr.Chatbot(show_copy_button=True, bubble_full_width=False, height=480)
-                user_message = gr.Textbox(label='User')
-
-                with gr.Row():
-                    user_message_btn = gr.Button('Send')
-                    stop_btn = gr.Button('Stop')
-                    clear_btn = gr.Button('Clear')
-
-                system_prompt = get_system_prompt_component(interactive=support_system_role.value)
-
-            with gr.Column(scale=1, min_width=80):
-                with gr.Group():
-                    gr.Markdown('Length of message history')
-                    history_len = gr.Slider(
-                        minimum=0,
-                        maximum=10,
-                        value=0,
-                        step=1,
-                        info='Number of previous messages taken into account in history',
-                        label='history_len',
-                        show_label=False,
-                        )
-
-                    with gr.Group():
-                        gr.Markdown('Generation parameters')
-                        do_sample = gr.Checkbox(
-                            value=False,
-                            label='do_sample',
-                            info='Activate random sampling',
-                            )
-                        generate_args = get_generate_args(do_sample.value)
-                        do_sample.change(
-                            fn=get_generate_args,
-                            inputs=do_sample,
-                            outputs=generate_args,
-                            show_progress=False,
-                        )
-
-        generate_event = gr.on(
-            triggers=[user_message.submit, user_message_btn.click],
-            fn=user_message_to_chatbot,
-            inputs=[user_message, chatbot],
-            outputs=[user_message, chatbot],
-        ).then(
-            fn=bot_response_to_chatbot,
-            inputs=[chatbot, model_dict, system_prompt, support_system_role, history_len, do_sample, *generate_args],
-            outputs=[chatbot],
-        )
-        stop_btn.click(
-            fn=None,
-            inputs=None,
-            outputs=None,
-            cancels=generate_event,
-        )
-        clear_btn.click(
-            fn=lambda: None,
-            inputs=None,
-            outputs=[chatbot],
-            )
-
-    # ================= LOAD MODELS PAGE ======================
-    with gr.Tab('Load model'):
-        gguf_url = gr.Textbox(
-            value='',
-            label='Link to GGUF',
-            placeholder='URL link to the model in GGUF format',
-            )
-        load_model_btn = gr.Button('Downloading GGUF and initializing the model')
-        load_log = gr.Textbox(
-            value=start_load_log,
-            label='Model loading status',
-            lines=3,
-            )
-            
-        load_model_btn.click(
-            fn=download_gguf_and_init_model,
-            inputs=[gguf_url, model_dict],
-            outputs=[model_dict, support_system_role, load_log],
-        ).success(
-            fn=get_system_prompt_component,
-            inputs=[support_system_role],
-            outputs=[system_prompt],
-        )
-
+from pathlib import Path
+from shutil import rmtree
+from typing import Union, List, Dict, Tuple, Optional
+from tqdm import tqdm
+
+import requests
+import gradio as gr
+from llama_cpp import Llama
+
+
+# ================== ANNOTATIONS ========================
+
+CHAT_HISTORY = List[Tuple[Optional[str], Optional[str]]]
+MODEL_DICT = Dict[str, Llama]
+
+
+# ================== FUNCS =============================
+
+def download_file(file_url: str, file_path: Union[str, Path]) -> None:
+    response = requests.get(file_url, stream=True)
+    if response.status_code != 200:
+        raise Exception(f'Файл недоступен для скачивания по ссылке: {file_url}')
+    total_size = int(response.headers.get('content-length', 0))
+    progress_tqdm = tqdm(desc='Loading GGUF file', total=total_size, unit='iB', unit_scale=True)
+    progress_gradio = gr.Progress()
+    completed_size = 0
+    with open(file_path, 'wb') as file:
+        for data in response.iter_content(chunk_size=4096):
+            size = file.write(data)
+            progress_tqdm.update(size)
+            completed_size += size
+            desc = f'Loading GGUF file, {completed_size/1024**3:.3f}/{total_size/1024**3:.3f} GB'
+            progress_gradio(completed_size/total_size, desc=desc)
+
+
+def download_gguf_and_init_model(gguf_url: str, model_dict: MODEL_DICT) -> Tuple[MODEL_DICT, bool, str]:
+    log = ''
+    if not gguf_url.endswith('.gguf'):
+        log += f'The link must be a direct link to the GGUF file\n'
+        return model_dict, log
+
+    gguf_filename = gguf_url.rsplit('/')[-1]
+    model_path = MODELS_PATH / gguf_filename
+    progress = gr.Progress()
+
+    if not model_path.is_file():
+        progress(0.3, desc='Шаг 1/2: Loading GGUF model file')
+        try:
+            download_file(gguf_url, model_path)
+            log += f'Model file {gguf_filename} successfully loaded\n'
+        except Exception as ex:
+            log += f'Error loading model from link {gguf_url}, error code:\n{ex}\n'
+            curr_model = model_dict.get('model')
+            if curr_model is None:
+                log += f'Model is missing from dictionary "model_dict"\n'
+                return model_dict, load_log
+            curr_model_filename = Path(curr_model.model_path).name
+            log += f'Current initialized model: {curr_model_filename}\n'
+            return model_dict, log
+    else:
+        log += f'Model file {gguf_filename} loaded, initializing model...\n'
+
+    progress(0.7, desc='Шаг 2/2: Model initialization')
+    model = Llama(model_path=str(model_path), n_gpu_layers=-1, verbose=True)
+    model_dict = {'model': model}
+    support_system_role = 'System role not supported' not in model.metadata['tokenizer.chat_template']
+    log += f'Model {gguf_filename} initialized\n'
+    return model_dict, support_system_role, log
+
+
+def user_message_to_chatbot(user_message: str, chatbot: CHAT_HISTORY) -> Tuple[str, CHAT_HISTORY]:
+    if user_message:
+        chatbot.append((user_message, None))
+    return '', chatbot
+
+
+def bot_response_to_chatbot(
+        chatbot: CHAT_HISTORY,
+        model_dict: MODEL_DICT,
+        system_prompt: str,
+        support_system_role: bool,
+        history_len: int,
+        do_sample: bool,
+        *generate_args,
+        ):
+
+    model = model_dict.get('model')
+    user_message = chatbot[-1][0]
+    messages = []
+
+    gen_kwargs = dict(zip(GENERATE_KWARGS.keys(), generate_args))
+    gen_kwargs['top_k'] = int(gen_kwargs['top_k'])
+
+    if not do_sample:
+        gen_kwargs['top_p'] = 0.0
+        gen_kwargs['top_k'] = 1
+        gen_kwargs['repeat_penalty'] = 1.0
+
+    if support_system_role and system_prompt:
+        messages.append({'role': 'system', 'content': system_prompt})
+
+    if history_len != 0:
+        for user_msg, bot_msg in chatbot[:-1][-history_len:]:
+            print(user_msg, bot_msg)
+            messages.append({'role': 'user', 'content': user_msg})
+            messages.append({'role': 'assistant', 'content': bot_msg})
+
+    messages.append({'role': 'user', 'content': user_message})
+    stream_response = model.create_chat_completion(
+        messages=messages,
+        stream=True,
+        **gen_kwargs,
+        )
+
+    chatbot[-1][1] = ''
+    for chunk in stream_response:
+        token = chunk['choices'][0]['delta'].get('content')
+        if token is not None:
+            chatbot[-1][1] += token
+            yield chatbot
+
+
+def get_system_prompt_component(interactive: bool) -> gr.Textbox:
+    value = '' if interactive else 'System prompt is not supported by this model'
+    return gr.Textbox(value=value, label='System prompt', interactive=interactive)
+
+
+def get_generate_args(do_sample: bool) -> List[gr.component]:
+    visible = do_sample
+    generate_args = [
+        gr.Slider(label='temperature', value=GENERATE_KWARGS['temperature'], minimum=0.1, maximum=3, step=0.1, visible=visible),
+        gr.Slider(label='top_p', value=GENERATE_KWARGS['top_p'], minimum=0.1, maximum=1, step=0.1, visible=visible),
+        gr.Slider(label='top_k', value=GENERATE_KWARGS['top_k'], minimum=1, maximum=50, step=5, visible=visible),
+        gr.Slider(label='repeat_penalty', value=GENERATE_KWARGS['repeat_penalty'], minimum=1, maximum=5, step=0.1, visible=visible),
+    ]
+    return generate_args
+
+
+# ================== VARIABLES =============================
+
+MODELS_PATH = Path('models')
+MODELS_PATH.mkdir(exist_ok=True)
+DEFAULT_GGUF_URL = 'https://huggingface.co/bartowski/gemma-2-2b-it-GGUF/resolve/main/gemma-2-2b-it-Q8_0.gguf'
+
+start_model_dict, start_support_system_role, start_load_log = download_gguf_and_init_model(
+    gguf_url=DEFAULT_GGUF_URL, model_dict={},
+    )
+
+GENERATE_KWARGS = dict(
+    temperature=0.2,
+    top_p=0.95,
+    top_k=40,
+    repeat_penalty=1.0,
+    )
+
+theme = gr.themes.Base(primary_hue='green', secondary_hue='yellow', neutral_hue='zinc').set(
+    loader_color='rgb(0, 255, 0)',
+    slider_color='rgb(0, 200, 0)',
+    body_text_color_dark='rgb(0, 200, 0)',
+    button_secondary_background_fill_dark='green',
+)
+css = '''.gradio-container {width: 60% !important}'''
+
+
+# ================== INTERFACE =============================
+
+with gr.Blocks(theme=theme, css=css) as interface:
+    model_dict = gr.State(start_model_dict)
+    support_system_role = gr.State(start_support_system_role)
+    
+    # ================= CHAT BOT PAGE ======================
+    with gr.Tab('Chat bot'):
+        with gr.Row():
+            with gr.Column(scale=3):
+                chatbot = gr.Chatbot(show_copy_button=True, bubble_full_width=False, height=480)
+                user_message = gr.Textbox(label='User')
+
+                with gr.Row():
+                    user_message_btn = gr.Button('Send')
+                    stop_btn = gr.Button('Stop')
+                    clear_btn = gr.Button('Clear')
+
+                system_prompt = get_system_prompt_component(interactive=support_system_role.value)
+
+            with gr.Column(scale=1, min_width=80):
+                with gr.Group():
+                    gr.Markdown('Length of message history')
+                    history_len = gr.Slider(
+                        minimum=0,
+                        maximum=10,
+                        value=0,
+                        step=1,
+                        info='Number of previous messages taken into account in history',
+                        label='history_len',
+                        show_label=False,
+                        )
+
+                    with gr.Group():
+                        gr.Markdown('Generation parameters')
+                        do_sample = gr.Checkbox(
+                            value=False,
+                            label='do_sample',
+                            info='Activate random sampling',
+                            )
+                        generate_args = get_generate_args(do_sample.value)
+                        do_sample.change(
+                            fn=get_generate_args,
+                            inputs=do_sample,
+                            outputs=generate_args,
+                            show_progress=False,
+                        )
+
+        generate_event = gr.on(
+            triggers=[user_message.submit, user_message_btn.click],
+            fn=user_message_to_chatbot,
+            inputs=[user_message, chatbot],
+            outputs=[user_message, chatbot],
+        ).then(
+            fn=bot_response_to_chatbot,
+            inputs=[chatbot, model_dict, system_prompt, support_system_role, history_len, do_sample, *generate_args],
+            outputs=[chatbot],
+        )
+        stop_btn.click(
+            fn=None,
+            inputs=None,
+            outputs=None,
+            cancels=generate_event,
+        )
+        clear_btn.click(
+            fn=lambda: None,
+            inputs=None,
+            outputs=[chatbot],
+            )
+
+    # ================= LOAD MODELS PAGE ======================
+    with gr.Tab('Load model'):
+        gguf_url = gr.Textbox(
+            value='',
+            label='Link to GGUF',
+            placeholder='URL link to the model in GGUF format',
+            )
+        load_model_btn = gr.Button('Downloading GGUF and initializing the model')
+        load_log = gr.Textbox(
+            value=start_load_log,
+            label='Model loading status',
+            lines=3,
+            )
+            
+        load_model_btn.click(
+            fn=download_gguf_and_init_model,
+            inputs=[gguf_url, model_dict],
+            outputs=[model_dict, support_system_role, load_log],
+        ).success(
+            fn=get_system_prompt_component,
+            inputs=[support_system_role],
+            outputs=[system_prompt],
+        )
+
+        gr.HTML("""<h3 style='text-align: center'>
+        <a href="https://github.com/sergey21000/gradio-llamacpp-chatbot" target='_blank'>GitHub Page</a></h3>
+        """)
+        
 interface.launch(server_name='0.0.0.0', server_port=7860)