import gradio as gr
import json
#from llmlingua import PromptCompressor
from utils_llmlingua2_test import PromptCompressor
import tiktoken

compressors = {
    "xlm-roberta": PromptCompressor(
        #model_name="microsoft/llmlingua-2-xlm-roberta-large-meetingbank",
        #model_name="qminh369/token-classification-llmlingua2-xlm-roberta-10k_merge_10_epoch_paper",
        #model_name='qminh369/token-classification-llmlingua2-xlm-roberta-42k_merge_1_epoch',
        model_name='qminh369/token-classification-llmlingua2-xlm-roberta-42k_merge_10_epoch',
        use_llmlingua2=True,
        device_map="cpu"
    )
}

tokenizer = tiktoken.encoding_for_model("gpt-4")

def compress(original_prompt, compression_rate, base_model="xlm-roberta", force_tokens = ['. ', ', '], chunk_end_tokens=['.', '\n']):
    if '\\n' in force_tokens:
        idx = force_tokens.index('\\n')
        force_tokens[idx] = '\n'

    compressor = compressors.get(base_model, compressors["xlm-roberta"])
    results = compressor.compress_prompt_llmlingua2(
            original_prompt,
            rate=compression_rate,
            force_tokens=force_tokens,
            chunk_end_tokens=chunk_end_tokens,
            return_word_label=True,
            drop_consecutive=True,
            force_reserve_digit=True,
            )

    compressed_prompt = results["compressed_prompt"]
    n_word_compressed = len(tokenizer.encode(compressed_prompt))
    
    word_sep = "\t\t|\t\t"
    label_sep = " "
    lines = results["fn_labeled_original_prompt"].split(word_sep)
    preserved_tokens = []
    for line in lines:
        word, label = line.split(label_sep)
        preserved_tokens.append((word, '+') if label == '1' else (word, None))

    return compressed_prompt, preserved_tokens, n_word_compressed

title = "LLMLingua-2"

header = """# LLMLingua-2
        """

theme = "soft"
css = """#anno-img .mask {opacity: 0.5; transition: all 0.2s ease-in-out;}
            #anno-img .mask.active {opacity: 0.7}"""

original_prompt_text = """"""

with gr.Blocks(title=title, css=css) as app:
    gr.Markdown(header)
    with gr.Row():
        with gr.Column(scale=3):
            original_prompt = gr.Textbox(value=original_prompt_text, label="Original Prompt", lines=10, max_lines=10, interactive=True)
            compressed_prompt = gr.Textbox(value='', label="Compressed Prompt", lines=10, max_lines=10, interactive=False)
            
        with gr.Column(scale=1):
            base_model = gr.Radio(["xlm-roberta"], label="Base Model", value="xlm-roberta", interactive=True)
            force_tokens = gr.Dropdown(['\\n', '.', '!', '?', ','],
                                       label="Tokens to Preserve",
                                       value=['\\n', '.', '!', '?', ','],
                                       multiselect=True,
                                       interactive=True)
            compression_rate = gr.Slider(minimum=0.1, maximum=1.0, step=0.1, value=0.7, label="Compression rate", info="after compr. / befor compr.", interactive=True)
            n_word_original = gr.Textbox(lines=1, label="Original (GPT-4 Tokens)", interactive=False, value=len(tokenizer.encode(original_prompt_text)))
            n_word_compressed = gr.Textbox(lines=1, label="Compressed (GPT-4 Tokens)", interactive=False)
    button = gr.Button("⚡Click to Compress")
    with gr.Accordion(label="Compression Details", open=False):
        diff_text = gr.HighlightedText(label="Diff", combine_adjacent=False, show_legend=True, color_map={"+": "green"})

    original_prompt.change(lambda x: len(tokenizer.encode(x)), inputs=[original_prompt], outputs=[n_word_original])
    original_prompt.change(lambda x: ("", "", []), inputs=[original_prompt], outputs=[compressed_prompt, n_word_compressed, diff_text])

    button.click(fn=compress,
                 inputs=[original_prompt, compression_rate, base_model, force_tokens],
                 outputs=[compressed_prompt, diff_text, n_word_compressed])
        
app.queue(max_size=10, api_open=False).launch(show_api=False)