Update app.py

app.py

@@ -0,0 +1,363 @@
+import os
+
+from transformers import AutoTokenizer
+from optimum.intel import OVModelForCausalLM, OVWeightQuantizationConfig
+from optimum.intel.openvino import OVModelForCausalLM
+from transformers import AutoConfig, AutoTokenizer
+import gradio as gr
+import time
+from threading import Thread
+
+from transformers import (
+    TextIteratorStreamer,
+    StoppingCriteria,
+    StoppingCriteriaList,
+    GenerationConfig,
+)
+# model_name = "openai-community/gpt2-large"
+# model_dir = "F:\\phi3\\openvinomodel\\phi3\\int4"
+# model_name = "savage1221/lora-fine"
+# save_name = model_name.split("/")[-1] + "_openvino"
+# precision = "f32"
+
+
+# quantization_config = OVWeightQuantizationConfig(
+#     bits=4,
+#     sym=False,
+#     group_size=128,
+#     ratio=0.6,
+#     trust_remote_code=True, 
+# )
+
+# ov_config = {"PERFORMANCE_HINT": "LATENCY", "NUM_STREAMS": "1", "CACHE_DIR": ""}
+
+# device = "gpu"
+
+
+# load_kwargs = {
+#     "device": device,
+#     "ov_config": {
+#         "PERFORMANCE_HINT": "LATENCY",
+#         # "INFERENCE_PRECISION_HINT": precision,
+#         "CACHE_DIR": os.path.join(save_name, "model_cache"),  # OpenVINO will use this directory as cache
+#     },
+#     "compile": False,
+#     "quantization_config": quantization_config,
+#     "trust_remote_code": True,
+#     # ov_config = ov_config
+# }
+
+# # Check whether the model was already exported
+# saved = os.path.exists(save_name)
+
+# model = OVModelForCausalLM.from_pretrained(
+#     # model_name
+#     model_name if not saved else save_name,
+#     export=not saved,
+#     **load_kwargs,
+# )
+# model = OVModelForCausalLM.from_pretrained(
+#     model_name,
+#     device='GPU.0',
+#     ov_config=ov_config,
+#     config=AutoConfig.from_pretrained(model_name, trust_remote_code=True),
+#     trust_remote_code=True,
+# )
+
+# # Load tokenizer to be used with the model
+# tokenizer = AutoTokenizer.from_pretrained(model_name if not saved else save_name)
+# tokenizer = AutoTokenizer.from_pretrained(model_name )
+
+# # Save the exported model locally
+# if not saved:
+#     model.save_pretrained(save_name)
+#     tokenizer.save_pretrained(save_name)
+
+# # TODO Optional: export to huggingface/hub
+
+# model_size = os.stat(os.path.join(save_name, "openvino_model.bin")).st_size / 1024 ** 3
+# print(f'Model size in FP32: ~5.4GB, current model size in 4bit: {model_size:.2f}GB')
+
+#####################################################################
+
+# Load model directly
+from transformers import AutoTokenizer, AutoModelForCausalLM
+
+tokenizer = AutoTokenizer.from_pretrained("savage1221/lora-fine", trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained("savage1221/lora-fine", trust_remote_code=True)
+
+
+# Copied and modified from https://github.com/bigcode-project/bigcode-evaluation-harness/blob/main/bigcode_eval/generation.py#L13
+class SuffixCriteria(StoppingCriteria):
+    def __init__(self, start_length, eof_strings, tokenizer, check_fn=None):
+        self.start_length = start_length
+        self.eof_strings = eof_strings
+        self.tokenizer = tokenizer
+        if check_fn is None:
+            check_fn = lambda decoded_generation: any(
+                [decoded_generation.endswith(stop_string) for stop_string in self.eof_strings]
+            )
+        self.check_fn = check_fn
+
+    def __call__(self, input_ids, scores, **kwargs):
+        """Returns True if generated sequence ends with any of the stop strings"""
+        decoded_generations = self.tokenizer.batch_decode(input_ids[:, self.start_length :])
+        return all([self.check_fn(decoded_generation) for decoded_generation in decoded_generations])
+
+
+def is_partial_stop(output, stop_str):
+    """Check whether the output contains a partial stop str."""
+    for i in range(0, min(len(output), len(stop_str))):
+        if stop_str.startswith(output[-i:]):
+            return True
+    return False
+
+
+
+# Set the chat template to the tokenizer. The chat template implements the simple template of
+#   User: content
+#   Assistant: content
+#   ...
+# Read more about chat templates here https://huggingface.co/docs/transformers/main/en/chat_templating
+tokenizer.chat_template = "{% for message in messages %}{{message['role'] + ': ' + message['content'] + '\n'}}{% endfor %}{% if add_generation_prompt %}{{ 'Assistant:' }}{% endif %}"
+
+
+def prepare_history_for_model(history):
+    """
+    Converts the history to a tokenized prompt in the format expected by the model.
+    Params:
+      history: dialogue history
+    Returns:
+      Tokenized prompt
+    """
+    messages = []
+    for idx, (user_msg, model_msg) in enumerate(history):
+        # skip the last assistant message if its empty, the tokenizer will do the formating
+        if idx == len(history) - 1 and not model_msg:
+            messages.append({"role": "User", "content": user_msg})
+            break
+        if user_msg:
+            messages.append({"role": "User", "content": user_msg})
+        if model_msg:
+            messages.append({"role": "Assistant", "content": model_msg})
+    input_token = tokenizer.apply_chat_template(
+        messages,
+        add_generation_prompt=True,
+        tokenize=True,
+        return_tensors="pt",
+        return_dict=True
+    )
+    return input_token
+
+
+def generate(history, temperature, max_new_tokens, top_p, repetition_penalty, assisted):
+    """
+    Generates the assistant's reponse given the chatbot history and generation parameters
+
+    Params:
+      history: conversation history formated in pairs of user and assistant messages `[user_message, assistant_message]`
+      temperature:  parameter for control the level of creativity in AI-generated text.
+                    By adjusting the `temperature`, you can influence the AI model's probability distribution, making the text more focused or diverse.
+      max_new_tokens: The maximum number of tokens we allow the model to generate as a response.
+      top_p: parameter for control the range of tokens considered by the AI model based on their cumulative probability.
+      repetition_penalty: parameter for penalizing tokens based on how frequently they occur in the text.
+      assisted: boolean parameter to enable/disable assisted generation with speculative decoding.
+    Yields:
+      Updated history and generation status.
+    """
+    start = time.perf_counter()
+    # Construct the input message string for the model by concatenating the current system message and conversation history
+    # Tokenize the messages string
+    inputs = prepare_history_for_model(history)
+    input_length = inputs['input_ids'].shape[1]
+    # truncate input in case it is too long.
+    # TODO improve this
+    if input_length > 2000:
+        history = [history[-1]]
+        inputs = prepare_history_for_model(history)
+        input_length = inputs['input_ids'].shape[1]
+
+    prompt_char = "▌"
+    history[-1][1] = prompt_char
+    yield history, "Status: Generating...", *([gr.update(interactive=False)] * 4)
+    
+    streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+
+    # Create a stopping criteria to prevent the model from playing the role of the user aswell.
+    stop_str = ["\nUser:", "\nAssistant:", "\nRules:", "\nQuestion:"]
+    stopping_criteria = StoppingCriteriaList([SuffixCriteria(input_length, stop_str, tokenizer)])
+    # Prepare input for generate
+    generation_config = GenerationConfig(
+        max_new_tokens=max_new_tokens,
+        do_sample=temperature > 0.0,
+        temperature=temperature if temperature > 0.0 else 1.0,
+        repetition_penalty=repetition_penalty,
+        top_p=top_p,
+        eos_token_id=[tokenizer.eos_token_id],
+        pad_token_id=tokenizer.eos_token_id,
+    )
+    generate_kwargs = dict(
+        streamer=streamer,
+        generation_config=generation_config,
+        stopping_criteria=stopping_criteria,
+    ) | inputs
+
+    if assisted:
+        target_generate = stateless_model.generate
+        generate_kwargs["assistant_model"] = asst_model
+    else:
+        target_generate = model.generate
+
+    t1 = Thread(target=target_generate, kwargs=generate_kwargs)
+    t1.start()
+
+    # Initialize an empty string to store the generated text.
+    partial_text = ""
+    for new_text in streamer:
+        partial_text += new_text
+        history[-1][1] = partial_text + prompt_char
+        for s in stop_str:
+            if (pos := partial_text.rfind(s)) != -1:
+                break
+        if pos != -1:
+            partial_text = partial_text[:pos]
+            break
+        elif any([is_partial_stop(partial_text, s) for s in stop_str]):
+            continue
+        yield history, "Status: Generating...", *([gr.update(interactive=False)] * 4)
+    history[-1][1] = partial_text
+    generation_time = time.perf_counter() - start
+    yield history, f'Generation time: {generation_time:.2f} sec', *([gr.update(interactive=True)] * 4)
+
+
+#############################################################
+
+
+# model.compile()
+
+
+try:
+    demo.close()
+except:
+    pass
+
+
+EXAMPLES = [
+    ["What is OpenVINO?"],
+    ["Can you explain to me briefly what is Python programming language?"],
+    ["Explain the plot of Cinderella in a sentence."],
+    ["Write a Python function to perform binary search over a sorted list. Use markdown to write code"],
+    ["Lily has a rubber ball that she drops from the top of a wall. The wall is 2 meters tall. How long will it take for the ball to reach the ground?"],
+]
+
+
+def add_user_text(message, history):
+    """
+    Add user's message to chatbot history
+
+    Params:
+      message: current user message
+      history: conversation history
+    Returns:
+      Updated history, clears user message and status
+    """
+    # Append current user message to history with a blank assistant message which will be generated by the model
+    history.append([message, None])
+    return ('', history)
+
+
+def prepare_for_regenerate(history):
+    """
+    Delete last assistant message to prepare for regeneration
+
+    Params:
+      history: conversation history
+    Returns:
+      updated history
+    """ 
+    history[-1][1] = None
+    return history
+
+
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown('<h1 style="text-align: center;">Chat with Phi-3 on Meteor Lake iGPU</h1>')
+    chatbot = gr.Chatbot()
+    with gr.Row():
+        assisted = gr.Checkbox(value=False, label="Assisted Generation", scale=10)
+        msg = gr.Textbox(placeholder="Enter message here...", show_label=False, autofocus=True, scale=75)
+        status = gr.Textbox("Status: Idle", show_label=False, max_lines=1, scale=15)
+    with gr.Row():
+        submit = gr.Button("Submit", variant='primary')
+        regenerate = gr.Button("Regenerate")
+        clear = gr.Button("Clear")
+    with gr.Accordion("Advanced Options:", open=False):
+        with gr.Row():
+            with gr.Column():
+                temperature = gr.Slider(
+                    label="Temperature",
+                    value=0.0,
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    interactive=True,
+                )
+                max_new_tokens = gr.Slider(
+                    label="Max new tokens",
+                    value=512,
+                    minimum=0,
+                    maximum=1024,
+                    step=32,
+                    interactive=True,
+                )
+            with gr.Column():
+                top_p = gr.Slider(
+                    label="Top-p (nucleus sampling)",
+                    value=1.0,
+                    minimum=0.0,
+                    maximum=1.0,
+                    step=0.05,
+                    interactive=True,
+                )
+                repetition_penalty = gr.Slider(
+                    label="Repetition penalty",
+                    value=1.0,
+                    minimum=1.0,
+                    maximum=2.0,
+                    step=0.1,
+                    interactive=True,
+                )
+    gr.Examples(
+        EXAMPLES, inputs=msg, label="Click on any example and press the 'Submit' button"
+    )
+
+    # Sets generate function to be triggered when the user submit a new message
+    gr.on(
+        triggers=[submit.click, msg.submit],
+        fn=add_user_text,
+        inputs=[msg, chatbot],
+        outputs=[msg, chatbot],
+        queue=False,
+    ).then(
+        fn=generate,
+        inputs=[chatbot, temperature, max_new_tokens, top_p, repetition_penalty, assisted],
+        outputs=[chatbot, status, msg, submit, regenerate, clear],
+        concurrency_limit=1,
+        queue=True
+    )
+    regenerate.click(
+        fn=prepare_for_regenerate,
+        inputs=chatbot,
+        outputs=chatbot,
+        queue=True,
+        concurrency_limit=1
+    ).then(
+        fn=generate,
+        inputs=[chatbot, temperature, max_new_tokens, top_p, repetition_penalty, assisted],
+        outputs=[chatbot, status, msg, submit, regenerate, clear],
+        concurrency_limit=1,
+        queue=True
+    )
+    clear.click(fn=lambda: (None, "Status: Idle"), inputs=None, outputs=[chatbot, status], queue=False)
+
+
+demo.launch()