Update app.py

app.py

@@ -1,214 +1,101 @@
 import gradio as gr
-from typing import Iterator, List, Tuple
-import torch
-from transformers import AutoModelForCausalLM, AutoTokenizer
-from peft import PeftConfig, PeftModel
-from huggingface_hub import login
-import os
-
-# Authenticate with Hugging Face
-HUGGINGFACE_TOKEN = os.getenv('HUGGINGFACE_TOKEN')  # Ensure this environment variable is set
-login(token=HUGGINGFACE_TOKEN)
-
-base_model = "mistralai/Mistral-7B-Instruct-v0.2"
-adapter = "GRMenon/mental-health-mistral-7b-instructv0.2-finetuned-V2"
-
-# Load tokenizer
-tokenizer = AutoTokenizer.from_pretrained(
-    base_model,
-    add_bos_token=True,
-    trust_remote_code=True,
-    padding_side='left'
-)
-
-# Create peft model using base_model and finetuned adapter
-config = PeftConfig.from_pretrained(adapter)
-model = AutoModelForCausalLM.from_pretrained(config.base_model_name_or_path,
-                                             load_in_4bit=True,
-                                             device_map='auto',
-                                             torch_dtype='auto')
-model = PeftModel.from_pretrained(model, adapter)
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-model.to(device)
-model.eval()
-
-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 = """
-# Simple Healthcare Chatbot
-### Powered by Mistral-7B with Healthcare Fine-Tuning
-"""
-
-def clear_and_save_textbox(message: str) -> tuple[str, str]:
-    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("Max new tokens exceeded")
-
-    history = history_with_input[:-1]
-    conversation = [{"role": "system", "content": system_prompt}] + \
-                   [{"role": "user", "content": user_input} for user_input, _ in history] + \
-                   [{"role": "user", "content": message}]
-    input_ids = tokenizer.apply_chat_template(conversation=conversation,
-                                              tokenize=True,
-                                              add_generation_prompt=True,
-                                              return_tensors='pt').to(device)
-    output_ids = model.generate(input_ids=input_ids, max_new_tokens=max_new_tokens,
-                                do_sample=True, pad_token_id=tokenizer.pad_token_id)
-    response = tokenizer.batch_decode(output_ids.detach().cpu().numpy(), skip_special_tokens=True)
-    response_text = response[0]
-
-    yield history + [(message, response_text)]
-
-def check_input_token_length(message: str, chat_history: list[tuple[str, str]], system_prompt: str) -> None:
-    input_token_length = len(tokenizer.encode(message)) + sum(len(tokenizer.encode(msg)) for msg, _ in 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:  # Link to CSS file
-    gr.Markdown(DESCRIPTION)
-    gr.Button("Duplicate Space for private use", elem_id="duplicate-button")
-
-    with gr.Group():
-        chatbot = gr.Chatbot(label="Chat with Healthcare AI")
-        with gr.Row():
-            textbox = gr.Textbox(
-                container=False,
-                show_label=False,
-                placeholder="Ask me anything about Healthcare and Wellness...",
-                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],
-    ).then(
-        fn=display_input,
-        inputs=[saved_input, chatbot],
-        outputs=chatbot,
-    ).then(
-        fn=check_input_token_length,
-        inputs=[saved_input, chatbot, system_prompt],
-    ).success(
-        fn=generate,
-        inputs=[saved_input, chatbot, system_prompt, max_new_tokens, temperature, top_p, top_k],
-        outputs=chatbot,
+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 ="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."
+
+# 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,
     )
 
-    submit_button.click(
-        fn=clear_and_save_textbox,
-        inputs=textbox,
-        outputs=[textbox, saved_input],
-    ).then(
-        fn=display_input,
-        inputs=[saved_input, chatbot],
-        outputs=chatbot,
-    ).then(
-        fn=check_input_token_length,
-        inputs=[saved_input, chatbot, system_prompt],
-    ).success(
-        fn=generate,
-        inputs=[saved_input, chatbot, system_prompt, max_new_tokens, temperature, top_p, top_k],
-        outputs=chatbot,
-    )
+    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
 
-    retry_button.click(
-        fn=delete_prev_fn,
-        inputs=chatbot,
-        outputs=[chatbot, saved_input],
-    ).then(
-        fn=display_input,
-        inputs=[saved_input, chatbot],
-        outputs=chatbot,
-    ).then(
-        fn=generate,
-        inputs=[saved_input, chatbot, system_prompt, max_new_tokens, temperature, top_p, top_k],
-        outputs=chatbot,
-    )
+def clear_fn():
+    return None, None
 
-    undo_button.click(
-        fn=delete_prev_fn,
-        inputs=chatbot,
-        outputs=[chatbot, saved_input],
-    ).then(
-        fn=lambda x: x,
-        inputs=[saved_input],
-        outputs=textbox,
-    )
+rand_val = random.randint(1, 1111111111111111)
 
-    clear_button.click(
-        fn=lambda: ([], ""),
-        outputs=[chatbot, saved_input],
-    )
+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))
 
-demo.queue(max_size=32).launch(share=False)
+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():
+        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)