from transformers import AutoTokenizer, AutoModelForCausalLM, StoppingCriteria
from peft import PeftModel
import torch
import gradio as gr
import os
import re

class ChineseCharacterStop(StoppingCriteria):
    def __init__(self, chars: list[str]):
        self.chars = [
            tokenizer(i, add_special_tokens=False, return_tensors='pt').input_ids
            for i in chars
        ]
        # for chars, tokens in zip(chars, self.chars):
        #     print(f"'{chars}':{tokens}")

    def __call__(self, input_ids: torch.LongTensor,
                 scores: torch.FloatTensor, **kwargs) -> bool:
        for c in self.chars:
            c = c.to(input_ids.device)
            match = torch.eq(input_ids[..., -c.shape[1]:], c)
            if torch.any(torch.all(match, dim=1)):
                return True
        return False


tokenizer = AutoTokenizer.from_pretrained("IDEA-CCNL/Wenzhong-GPT2-110M")
tokenizer.pad_token = tokenizer.eos_token
gpt2_model = AutoModelForCausalLM.from_pretrained("IDEA-CCNL/Wenzhong-GPT2-110M")
model = PeftModel.from_pretrained(gpt2_model, 'checkpoint_lora_v4.1')


def cang_tou(tou: str):
    poem_now = "写一首唐诗："
    for c in tou:
        poem_now += c
        print(poem_now)
        inputs = tokenizer(poem_now, return_tensors='pt')
        outputs = model.generate(
            **inputs,
            return_dict_in_generate=True,
            max_length=150,
            do_sample=True,
            top_p=0.4,
            num_beams=1,
            num_return_sequences=1,
            stopping_criteria=[ChineseCharacterStop(['。', '，'])],
            pad_token_id=tokenizer.pad_token_id
        )
        poem_now = tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True)[0]
    print(poem_now)
    return poem_now[6:]


def prompt_gen(prompt):
    inputs = tokenizer(prompt, return_tensors='pt')
    outputs = model.generate(
        **inputs,
        return_dict_in_generate=True,
        max_length=200,
        do_sample=True,
        top_p=0.8,
        num_beams=5,
        num_return_sequences=3,
        # stopping_criteria=[ChineseCharacterStop(['。', '，', ''])],
        pad_token_id=tokenizer.pad_token_id
    )
    res = ''
    for line in tokenizer.batch_decode(outputs.sequences, skip_special_tokens=True):
      line = line[len(prompt):]
      res = res+line+'\n'
    return res

css = """
#col-container {max-width: 510px; margin-left: auto; margin-right: auto;}
a {text-decoration-line: underline; font-weight: 600;}
.animate-spin {
  animation: spin 1s linear infinite;
}
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown(
            """
            <h1 style="text-align: center;">✨古诗生成</h1>
            <p style="text-align: center;">
            根据输入的提示生成古诗、藏头诗<br />
            </p>            
            """
        )
        with gr.Tab("提示"):
            prompt_in = gr.Textbox(label="Prompt", placeholder="写一首关于思乡的古诗：", elem_id="prompt-in")
            #neg_prompt = gr.Textbox(label="Negative prompt", value="text, watermark, copyright, blurry, nsfw", elem_id="neg-prompt-in")
            #inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, step=1, value=40, interactive=False)
            submit_btn = gr.Button("Submit")
            poetry_result = gr.Textbox(label="Output", elem_id="poetry-output")

            submit_btn.click(fn=prompt_gen,
                    inputs=[prompt_in],
                    outputs=[poetry_result])
        
        with gr.Tab("藏头诗"):
            tou_in = gr.Textbox(label="Prompt", placeholder="一见如故", elem_id="tou-in")
            #neg_prompt = gr.Textbox(label="Negative prompt", value="text, watermark, copyright, blurry, nsfw", elem_id="neg-prompt-in")
            #inference_steps = gr.Slider(label="Inference Steps", minimum=10, maximum=100, step=1, value=40, interactive=False)
            submit_btn = gr.Button("Submit")
            cangtou_result = gr.Textbox(label="Output", elem_id="cangtou-output")
            submit_btn.click(fn=cang_tou,
                    inputs=[tou_in],
                    outputs=[cangtou_result])
        
    

demo.queue(max_size=12).launch()