import gradio as gr
from compressor import PromptCompressor

import os
res = os.popen('python -m spacy download en_core_web_sm').read()
print(res)


def compressit(original_text, compressor1, ratio, maxlength):

    if compressor1=="Selective Context":
        compressor = PromptCompressor(type='SCCompressor', lang='en', model='gpt2', device='cpu')
    elif compressor1=="LLMLingua":
        return "Sorry, currently we cannot provide services for LLMLingua due to the Huggingface Token issue. Please try other compressors."
    elif compressor1=="LongLLMLingua":
        return "Sorry, currently we cannot provide services for LongLLMLingua due to the Huggingface Token issue. Please try other compressors."
    elif compressor1=="SCRL":
        compressor = PromptCompressor(type='SCRLCompressor', model_dir="models/gigaword-L8/", device="cpu", tokenizer_dir="sentence-transformers/paraphrase-distilroberta-base-v2")
    elif compressor1=="KiS":
        compressor = PromptCompressor(type='KiSCompressor', device="cuda", model_dir="philippelaban/keep_it_simple")
    else:
        compressor = PromptCompressor(type='SCCompressor', lang='en', model='gpt2', device='cpu')

    if compressor1 != "SCRL":
        compressed_prompt = compressor.compressgo(original_prompt=original_text, ratio=float(ratio), max_length=int(maxlength))
    else:
        compressed_prompt = compressor.compressgo(original_prompt=original_text, ratio=float(ratio), max_length=int(maxlength))
    return compressed_prompt["compressed_prompt"]


demo = gr.Interface(
    fn=compressit,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your prompt here...", label="input", info="Enter the original prompt here."),
        gr.Dropdown(
            ["Selective Context", "LLMLingua", "LongLLMLingua", "SCRL", "KiS"], label="compressor", info="Choose your compressor here. \n Currently, we cannot support the online demo for LLMLingua and LongLLMLingua due to the Huggingface Token issue."
        ),
        gr.Textbox(lines=1, placeholder="Enter the compression ratio here...", info="Ratio only works for Selective Context, LLMLingua and LongLLMLingua."),
        gr.Textbox(lines=1, placeholder="Enter the max_length parameter if you are using SCRL or KiS", label="max_length", info="If you are using SCRL or KiS, fill in the parameter, if not, just ignore this.\n Hint: For SCRL, max_length should be shorter than the lenght of original prompt; For KiS, max_length should be longer than it.")
    ],
    outputs=[
        gr.Textbox(lines=1, info="Please note that when the text is very short, LLMLingua and LongLLMLingua will not work.")
    ]
)

demo.launch(share=False)