#from summary_reverse_pred_native import *
#### daspartho/prompt-extend

import os 
os.system("pip install huggingface_hub")

from huggingface_hub import space_info

import gradio as gr
#import os
from predict import *

#device = "cuda:0"
device = "cpu"
assert device.startswith("cpu") or device.startswith("cuda")

from transformers import (
    T5ForConditionalGeneration,
    MT5ForConditionalGeneration,
    ByT5Tokenizer,
    PreTrainedTokenizer,
    T5TokenizerFast as T5Tokenizer,
    MT5TokenizerFast as MT5Tokenizer,
    AutoModelForSeq2SeqLM,
    AutoTokenizer,
    BertTokenizer,
    GPT2LMHeadModel,
)

#### "svjack/prompt-extend-chinese-gpt"
#model_path = "/home/featurize/zh_p_extend_outputs/simplet5-epoch-3-train-loss-1.2628-val-loss-1.6293"
model_path = "svjack/prompt-extend-chinese-gpt"
tokenizer1 = BertTokenizer.from_pretrained(model_path)
model1 = GPT2LMHeadModel.from_pretrained(model_path)

if device.startswith("cuda"):
    zh_pe_model = Obj(model1, tokenizer1, device = "cuda:0")
else:
    zh_pe_model = Obj(model1, tokenizer1, device = "cpu")

def one_ele_trans(x):
    x = x.strip()
    x = x[1:] if x.startswith("'") else x
    x = x[:-1] if x.endswith("'") else x
    x = x[1:] if x.startswith('"') else x
    x = x[:-1] if x.endswith('"') else x
    return x

def stdf_prompt_expander(x, do_sample):
    assert type(x) == type("")
    return zh_pe_model.predict(
    one_ele_trans(x.strip()).strip(),
    max_length = 128,
    do_sample = do_sample
    )[0].replace(" ", "").strip()

#text0 = "飓风格特是1993年9月在墨西哥和整个中美洲引发严重洪灾的大规模热带气旋，源于9月14日西南加勒比海上空一股东风波。次日从尼加拉瓜登岸，经过洪都拉斯后于9月17日在洪都拉斯湾再次达到热带风暴标准，但次日进入伯利兹上空后就减弱成热带低气压。穿过尤卡坦半岛后，在9月20日强化成二级飓风，从韦拉克鲁斯州的图斯潘附近登陆墨西哥。9月21日从纳亚里特州进入太平洋时已降级成热带低气压，最终于5天后在开放水域上空消散。"
#text1 = "珊瑚坝是长江中的一处河漫滩，位于长江重庆市渝中区区段主航道左侧[1]，靠近渝中半岛，原分属重庆市市中区菜园坝街道和石板坡街道[2]，现属渝中区菜园坝街道石板坡社区[3]，是长江上游缓冲地段自然冲积沙洲，略呈纺锤形[4]或椭圆形，长约1800米，宽约600米，坝上遍布鹅卵石和水草。每年夏季洪水时均被淹没，其余时间常露水面，枯水期则与长江左岸相连[5]。"
prompt = "一只凶猛的老虎，咬死了一只豺狼。"

example_sample = [
    [prompt, False],
    #[text1, False],
]

markdown_exp_size = "##"
lora_repo = "svjack/chatglm3-few-shot"
lora_repo_link = "svjack/chatglm3-few-shot/?input_list_index=9"
emoji_info = space_info(lora_repo).__dict__["cardData"]["emoji"]
space_cnt = 1
task_name = "[---Stable Diffusion Chinese Prompt Extend---]"
description = f"{markdown_exp_size} {task_name} few shot prompt in ChatGLM3 Few Shot space repo (click submit to activate) : [{lora_repo_link}](https://huggingface.co/spaces/{lora_repo_link}) {emoji_info}"


def demo_func(prefix, do_sample):
    #l = simple_pred(prefix, do_sample = do_sample)
    x = stdf_prompt_expander(prefix, do_sample = do_sample)
    return {
        "Prompt extend": x
    }

demo = gr.Interface(
        fn=demo_func,
        inputs=[gr.Text(label = "Prompt"),
                gr.Checkbox(label="do sample"),
        ],
        outputs="json",
        title=f"Stable Diffusion Chinese Prompt Extend 🐰 demonstration",
        description = 'This _example_ was **drive** from <br/><b><h4>[https://github.com/svjack/Stable-Diffusion-Chinese-Extend](https://github.com/svjack/Stable-Diffusion-Chinese-Extend)</h4></b>\n',
        #description = description,
        examples=example_sample if example_sample else None,
        cache_examples = False
    )


with demo:
    gr.HTML(
    '''
                <div style="justify-content: center; display: flex;">
                <iframe
                src="https://svjack-chatglm3-few-shot-demo.hf.space/?input_list_index=9"
                frameborder="0"
                width="1400"
                height="768"
                ></iframe>
                </div>
    '''
    )


demo.launch(server_name=None, server_port=None)