Update app.py

app.py

@@ -1,101 +1,305 @@
+import os
+from typing import Iterator
+
 import gradio as gr
-from huggingface_hub import InferenceClient
-import random
-import textwrap
-
-# Define the model to be used
-model = "mistralai/Mixtral-8x7B-Instruct-v0.1"
-
-# Load model directly
-#model = "GRMenon/mental-health-mistral-7b-instructv0.2-finetuned-V2"
-client = InferenceClient(model)
-
-# Embedded system prompt
-system_prompt_text = "Act like a compassionate and helpful Health consultant and professional named PhoenixA. You help and support with any kind of request and provide a detailed answer or suggestion to the question. You are friendly and always willing to help. But if you are asked about something unethical or dangerous, you must provide a safe and respectful way to handle that. Never say that you cannot help themBe sure to ask for specific problem and do your best to give professional advices, remember you are a professional."
-
-# Read the content of the info.md file
-with open("info.md", "r") as file:
-    info_md_content = file.read()
-
-# Chunk the info.md content into smaller sections
-chunk_size = 2000  # Adjust this size as needed
-info_md_chunks = textwrap.wrap(info_md_content, chunk_size)
-
-def get_all_chunks(chunks):
-    return "\n\n".join(chunks)
-
-def format_prompt_mixtral(message, history, info_md_chunks):
-    prompt = "<s>"
-    all_chunks = get_all_chunks(info_md_chunks)
-    prompt += f"{all_chunks}\n\n"  # Add all chunks of info.md at the beginning
-    prompt += f"{system_prompt_text}\n\n"  # Add the system prompt
-
-    if history:
-        for user_prompt, bot_response in history:
-            prompt += f"[INST] {user_prompt} [/INST]"
-            prompt += f" {bot_response}</s> "
-    prompt += f"[INST] {message} [/INST]"
-    return prompt
-
-def chat_inf(prompt, history, seed, temp, tokens, top_p, rep_p):
-    generate_kwargs = dict(
-        temperature=temp,
-        max_new_tokens=tokens,
-        top_p=top_p,
-        repetition_penalty=rep_p,
-        do_sample=True,
-        seed=seed,
+
+from src.model import run
+
+HF_PUBLIC = os.environ.get("HF_PUBLIC", False)
+
+DEFAULT_SYSTEM_PROMPT = "You are Phoenix AI Healthcare. You are professional, you are polite, give only truthful information and are based on the Mistral-7B model from Mistral AI about Healtcare and Wellness. You can communicate in different languages equally well."
+MAX_MAX_NEW_TOKENS = 4096
+DEFAULT_MAX_NEW_TOKENS = 256
+MAX_INPUT_TOKEN_LENGTH = 4000
+
+DESCRIPTION = """
+# [Mistral-7B](https://huggingface.co/mistralai/Mistral-7B-Instruct-v0.1)
+"""
+
+
+def clear_and_save_textbox(message: str) -> tuple[str, str]:
+    """
+    Clear the textbox and save the input to a state variable.
+    :param message: The input message.
+    :return: A tuple of the empty string and the input message.
+    """
+    return "", message
+
+
+def display_input(
+    message: str, history: list[tuple[str, str]]
+) -> list[tuple[str, str]]:
+    """
+    Display the input message in the chat history.
+    :param message: The input message.
+    :param history: The chat history.
+    :return: The chat history with the input message appended.
+    """
+    history.append((message, ""))
+    return history
+
+
+def delete_prev_fn(
+        history: list[tuple[str, str]]) -> tuple[list[tuple[str, str]], str]:
+    """
+    Delete the previous message from the chat history.
+    :param history: The chat history.
+    :return: The chat history with the last message removed
+    and the removed message.
+    """
+    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]]]:
+    """
+    Generate a response to the input message.
+    :param message: The input message.
+    :param history_with_input: The chat history with
+    the input message appended.
+    :param system_prompt: The system prompt.
+    :param max_new_tokens: The maximum number of tokens to generate.
+    :param temperature: The temperature.
+    :param top_p: The top-p (nucleus sampling) probability.
+    :param top_k: The top-k probability.
+    :return: An iterator over the chat history with
+    the generated response appended.
+    """
+    if max_new_tokens > MAX_MAX_NEW_TOKENS:
+        raise ValueError
+
+    history = history_with_input[:-1]
+    generator = 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)]
 
-    formatted_prompt = format_prompt_mixtral(prompt, history, info_md_chunks)
-    stream = client.text_generation(formatted_prompt, **generate_kwargs, stream=True, details=True, return_full_text=False)
-    output = ""
-    for response in stream:
-        output += response.token.text
-        yield [(prompt, output)]
-    history.append((prompt, output))
-    yield history
 
-def clear_fn():
-    return None, None
+def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:
+    """
+    Process an example.
+    :param message: The input message.
+    :return: A tuple of the empty string and the chat history with the \
+        generated response appended.
+    """
+    generator = generate(message, [], DEFAULT_SYSTEM_PROMPT, 1024, 1, 0.95, 50)
+    for x in generator:
+        pass
+    return "", x
 
-rand_val = random.randint(1, 1111111111111111)
 
-def check_rand(inp, val):
-    if inp:
-        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=random.randint(1, 1111111111111111))
-    else:
-        return gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, value=int(val))
+def check_input_token_length(
+    message: str, chat_history: list[tuple[str, str]], system_prompt: str
+) -> None:
+    """
+    Check that the accumulated input is not too long.
+    :param message: The input message.
+    :param chat_history: The chat history.
+    :param system_prompt: The system prompt.
+    :return: None.
+    """
+    input_token_length = len(message) + len(chat_history)
+    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."
+        )
+
+
+with gr.Blocks(css="./styles/style.css") as demo:
+    gr.Markdown(DESCRIPTION)
+    gr.DuplicateButton(
+        value="Duplicate Space for private use", elem_id="duplicate-button"
+    )
 
-with gr.Blocks() as app:  # Add auth here
-    gr.HTML("""<center><h1 style='font-size:xx-large;'>PhoenixAI</h1><br><h3> made with love by Omdena </h3><br><h7>EXPERIMENTAL</center>""")
-    with gr.Row():
-        chat = gr.Chatbot(height=500)
     with gr.Group():
+        chatbot = gr.Chatbot(label="Playground")
         with gr.Row():
-            with gr.Column(scale=3):
-                inp = gr.Textbox(label="Prompt", lines=5, interactive=True)  # Increased lines and interactive
-                with gr.Row():
-                    with gr.Column(scale=2):
-                        btn = gr.Button("Chat")
-                    with gr.Column(scale=1):
-                        with gr.Group():
-                            stop_btn = gr.Button("Stop")
-                            clear_btn = gr.Button("Clear")
-            with gr.Column(scale=1):
-                with gr.Group():
-                    rand = gr.Checkbox(label="Random Seed", value=True)
-                    seed = gr.Slider(label="Seed", minimum=1, maximum=1111111111111111, step=1, value=rand_val)
-                    tokens = gr.Slider(label="Max new tokens", value=3840, minimum=0, maximum=8000, step=64, interactive=True, visible=True, info="The maximum number of tokens")
-                    temp = gr.Slider(label="Temperature", step=0.01, minimum=0.01, maximum=1.0, value=0.9)
-                    top_p = gr.Slider(label="Top-P", step=0.01, minimum=0.01, maximum=1.0, value=0.9)
-                    rep_p = gr.Slider(label="Repetition Penalty", step=0.1, minimum=0.1, maximum=2.0, value=1.0)
-
-    hid1 = gr.Number(value=1, visible=False)
-
-    go = btn.click(check_rand, [rand, seed], seed).then(chat_inf, [inp, chat, seed, temp, tokens, top_p, rep_p], chat)
-
-    stop_btn.click(None, None, None, cancels=[go])
-    clear_btn.click(clear_fn, None, [inp, chat])
-
-app.queue(default_concurrency_limit=10).launch(share=True)
+            textbox = gr.Textbox(
+                container=False,
+                show_label=False,
+                placeholder="Greetings, with what Healthcare/Wellness topic can I help you with today?",
+                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.Accordion(label="⚙️ Advanced options", open=False):
+        system_prompt = gr.Textbox(
+            label="System prompt",
+            value=DEFAULT_SYSTEM_PROMPT,
+            lines=5,
+            interactive=False,
+        )
+        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=0.1,
+        )
+        top_p = gr.Slider(
+            label="Top-p (nucleus sampling)",
+            minimum=0.05,
+            maximum=1.0,
+            step=0.05,
+            value=0.9,
+        )
+        top_k = gr.Slider(
+            label="Top-k",
+            minimum=1,
+            maximum=1000,
+            step=1,
+            value=10,
+        )
+
+    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=32).launch(share=HF_PUBLIC, show_api=False)