import gradio as gr
from llmlingua import PromptCompressor

llm_lingua = PromptCompressor("lgaalves/gpt2-dolly", device_map="cpu")

INTRO = """
# LLMLingua: Compressing Prompts for Accelerated Inference of Large Language Models
This is an early demo of the prompt compression method LLMLingua.
To use it, upload your prompt and set the compression target.
1. ✅ Set the different components of the prompt separately, including instruction, context, and question. Leave the corresponding field empty if a particular component does not exist.
    - Question: This refers to the directives given by the user to the LLMs, such as inquiries, questions, or requests. Positioned after the instruction and context modules, the question module has a high sensitivity to compression.
    - Context: This module provides the supplementary context needed to address the question, such as documents, demonstrations, web search results, or API call results. Located between the instruction and question modules, its sensitivity to compression is relatively low.
    - Instruction: This module consists of directives given by the user to the LLMs, such as task descriptions. Placed before the instruction and context modules, the instruction module exhibits a high sensitivity to compression.
2. ✅ Set the target_token or compression ratio.
3. 🤔 Try experimenting with different target compression ratios or other hyperparameters to optimize the performance.

You can check our [repo](https://aka.ms/LLMLingua)!
"""

custom_css = """
    #image-upload {
        flex-grow: 1;
    }
    #params .tabs {
        display: flex;
        flex-direction: column;
        flex-grow: 1;
    }
    #params .tabitem[style="display: block;"] {
        flex-grow: 1;
        display: flex !important;
    }
    #params .gap {
        flex-grow: 1;
    }
    #params .form {
        flex-grow: 1 !important;
    }
    #params .form > :last-child{
        flex-grow: 1;
    }
    .md ol, .md ul {
        margin-left: 1rem;
    }
    .md img {
        margin-bottom: 1rem;
    }
"""

def compress_prompt(context, instruction, question, ratio, target_token):
    context, instruction, question = context.replace("\\n", "\n"), instruction.replace("\\n", "\n"), question.replace("\\n", "\n")
    compressed_prompt = llm_lingua.compress_prompt(context.split("\n\n"), instruction, question, float(ratio), float(target_token))

    return [compressed_prompt[key] for key in ["compressed_prompt", "origin_tokens", "compressed_tokens", "ratio", "saving"]]


with gr.Blocks(css=custom_css) as iface:
    gr.Markdown(INTRO)

    with gr.Row():
        with gr.Column(elem_id="params"):
            with gr.Tab('Prompts'):
                instruction = gr.Textbox(
                    label="Instruction",
                    lines=1,
                    value="",
                )
                context = gr.Textbox(
                    label="Context",
                    lines=3,
                    value="",
                )
                question = gr.Textbox(
                    label="Question",
                    lines=1,
                    value="",
                )
            with gr.Tab('Compression Target'):
                ratio = gr.Textbox(
                    label="Compression Ratio",
                    value=0.5,
                )
                target_token = gr.Textbox(
                    label="Target Token",
                    value=-1,
                )

    gen_button = gr.Button(value="Compress Prompt!", variant="primary")

    # with gr.Tab('Results'):
    #     results = gr.Gallery(
    #         show_label=False,
    #         object_fit="contain",
    #         columns=4
    #     )

    with gr.Row():
        with gr.Column(elem_id="Results"):
            with gr.Tab('Compressed Prompts'):
                compressed_prompt = gr.Textbox(
                    label="compressed_prompt",
                    lines=5,
                )
            with gr.Tab('Saving'):
                origin_tokens = gr.Textbox(
                    label="The tokens number of original prompt",
                )
                compressed_tokens = gr.Textbox(
                    label="The tokens number of compressed prompt",
                )
                saving_ratio = gr.Textbox(
                    label="Actual Compression Ratio",
                )
                saving = gr.Textbox(
                    label="Saving Cost",
                )


    # gr.Examples(
    #     examples=EXAMPLES,
    #     inputs=[image_upload, positive_prompt, negative_prompt],
    # )

    gen_button.click(
        fn=compress_prompt,
        inputs=[
            context,
            instruction,
            question,
            ratio,
            target_token
        ],
        outputs=[
            compressed_prompt,
            origin_tokens,
            compressed_tokens,
            saving_ratio,
            saving
        ],
    )

iface.queue(max_size=10, api_open=False).launch(show_api=False)