import os
import io
import json
import random
import string
import requests
import base64
import torch

import gradio as gr
from PIL import Image
from datasets import load_dataset
from datasets import DownloadMode, VerificationMode
from diffusers import StableDiffusionXLPipeline

from pingpong import PingPong
from pingpong.pingpong import PPManager
from pingpong.pingpong import PromptFmt
from pingpong.pingpong import UIFmt
from pingpong.gradio import GradioChatUIFmt

from fpdf import FPDF

class PDF(FPDF):
    def header(self):
        # Arial bold 15
        self.set_font('Arial', 'B', 15)
        # Calculate width of title and position
        w = self.get_string_width(self.title) + 6
        self.set_x((210 - w) / 2)
        # Colors of frame, background and text
        self.set_draw_color(255, 255, 255)
        self.set_fill_color(255, 255, 255)
        # self.set_text_color(220, 50, 50)
        # Thickness of frame (1 mm)
        self.set_line_width(1)
        # Title
        self.cell(w, 9, self.title, 1, 1, 'C', 1)
        # Line break
        self.ln(10)

        if self.art is not None:
          self.image(self.art, x=self.w/2.0-25, w=50)
          self.ln(10)

    def footer(self):
        # Position at 1.5 cm from bottom
        self.set_y(-15)
        # Arial italic 8
        self.set_font('Arial', 'I', 8)
        # Text color in gray
        self.set_text_color(128)
        # Page number
        self.cell(0, 10, 'Page ' + str(self.page_no()), 0, 0, 'C')

    def chapter_title(self, num, label):
        # Arial 12
        self.set_font('Arial', '', 12)
        # Background color
        self.set_fill_color(200, 220, 255)
        # Title
        self.cell(0, 6, 'Chapter %d : %s' % (num, label), 0, 1, 'L', 1)
        # Line break
        self.ln(4)

    def chapter_body(self, content):
        # Times 12
        self.set_font('Times', '', 12)
        # Output justified text
        self.multi_cell(0, 5, content)
        # Line break
        self.ln()
        # Mention in italics
        self.set_font('', 'I')

    def print_chapter(self, content):
        self.add_page()
        self.chapter_body(content)

class LLaMA2ChatPromptFmt(PromptFmt):
    @classmethod
    def ctx(cls, context):
        if context is None or context == "":
            return ""
        else:
            return f"""<<SYS>>
{context}
<</SYS>>

"""

    @classmethod
    def prompt(cls, pingpong, truncate_size):
        ping = pingpong.ping[:truncate_size]
        pong = "" if pingpong.pong is None else pingpong.pong[:truncate_size]
        return f"""[INST] {ping} [/INST] {pong}"""

class LLaMA2ChatPPManager(PPManager):
    def build_prompts(self, from_idx: int=0, to_idx: int=-1, fmt: PromptFmt=LLaMA2ChatPromptFmt, truncate_size: int=None):
        if to_idx == -1 or to_idx >= len(self.pingpongs):
            to_idx = len(self.pingpongs)

        results = fmt.ctx(self.ctx)

        for idx, pingpong in enumerate(self.pingpongs[from_idx:to_idx]):
            results += fmt.prompt(pingpong, truncate_size=truncate_size)

        return results

class GradioLLaMA2ChatPPManager(LLaMA2ChatPPManager):
    def build_uis(self, from_idx: int=0, to_idx: int=-1, fmt: UIFmt=GradioChatUIFmt):
        if to_idx == -1 or to_idx >= len(self.pingpongs):
            to_idx = len(self.pingpongs)

        results = []

        for pingpong in self.pingpongs[from_idx:to_idx]:
            results.append(fmt.ui(pingpong))

        return results

TOKEN = os.getenv('HF_TOKEN')
MODEL_ID = 'meta-llama/Llama-2-70b-chat-hf'

pipe = StableDiffusionXLPipeline.from_pretrained(
    "stabilityai/stable-diffusion-xl-base-1.0", 
    torch_dtype=torch.float16, 
    variant="fp16", 
    use_safetensors=True
)
pipe.to("cuda")

STYLES = """
.margin-auto {
  margin: auto !important;
}

.left-panel {
  min-width: min(290px, 100%) !important;
}

.small-big {
  font-size: 12pt !important;
}

.small-big-textarea > label > textarea {
  font-size: 12pt !important;
}

.highlighted-text {
  background: yellow;
  overflow-wrap: break-word;
}

.no-gap {
  gap: 0px !important;
}

.group-border {
  padding: 10px;
  border-width: 1px;
  border-radius: 10px;
  border-color: gray;
  border-style: dashed;
}

.control-label-font {
  font-size: 13pt !important;
}

.control-button {
  background: none !important;
  border-color: #69ade2 !important;
  border-width: 2px !important;
  color: #69ade2 !important;
}

.center {
  text-align: center;
}

.right {
  text-align: right;
}

.no-label {
  padding: 0px !important;
}

.no-label > label > span {
  display: none;
}

.no-label-chatbot {
  border: none !important;
  box-shadow: none !important;
  height: 520px !important;
}

.no-label-chatbot > div > div:nth-child(1) {
  display: none;
}

.no-label-image > div:nth-child(2) {
  display: none;
}

.left-margin-30 {
  padding-left: 30px !important;
}

.left {
  text-align: left !important;
}

.alt-button {
  color: gray !important;
  border-width: 1px !important;
  background: none !important;
  border-color: gray !important;
  text-align: justify !important;
}

.white-text {
  color: #000 !important;
}
"""

def id_generator(size=6, chars=string.ascii_uppercase + string.digits):
  return ''.join(random.choice(chars) for _ in range(size))

def get_new_ppm(ping):
  ppm = LLaMA2ChatPPManager()
  ppm.ctx = """\
  You are a helpful, respectful and honest writing helper. Always write stories that suites to query.

  You DO NOT give explanation but just stories. For instance, do not say such as "Sure! Here's a short paragraph to start a short story:"""

  ppm.add_pingpong(PingPong(ping, ''))
  return ppm

def get_new_ppm_for_chat():
  ppm = GradioLLaMA2ChatPPManager()
  return ppm

def gen_text(prompt, hf_model='meta-llama/Llama-2-70b-chat-hf', hf_token=None, parameters=None):
  if hf_token is None:
    raise ValueError("Hugging Face Token is not set")

  if parameters is None:
    parameters = {
        'max_new_tokens': 512,
        'do_sample': True,
        'return_full_text': False,
        'temperature': 1.0,
        'top_k': 50,
        # 'top_p': 1.0,
        'repetition_penalty': 1.2
    }

  url = f'https://api-inference.huggingface.co/models/{hf_model}'
  headers={
      'Authorization': f'Bearer {hf_token}',
      'Content-type': 'application/json'
  }
  data = {
      'inputs': prompt,
      'stream': False,
      'options': {
          'use_cache': False,
      },
      'parameters': parameters
  }

  r = requests.post(
      url,
      headers=headers,
      data=json.dumps(data)
  )

  if r.reason != 'OK':
    raise ValueError("Response other than 200")

  return json.loads(r.content.decode("utf-8"))[0]['generated_text']

def gen_art(editor, cover_art_image, gen_cover_art_prompt):
  if gen_cover_art_prompt.strip() == "":
    ppm = get_new_ppm(f"""describe the story below as a movie poster with ONLY general terms and adjectives. give me the caption ONLY.
--------------------------------
{editor}""")
    cover_art_prompt = gen_text(ppm.build_prompts(), hf_model=MODEL_ID, hf_token=TOKEN)

    return [
        cover_art_image,
        cover_art_prompt
    ]
  else:
    global pipe
    pipe = pipe.to("cuda")

    return [
        pipe(gen_cover_art_prompt, num_inference_steps=25).images[0].resize((512, 512)),
        gen_cover_art_prompt
    ]

def generate_pdf(title, editor, concept_art):
  tmp_filename = id_generator()

  if concept_art is not None:
    im = Image.fromarray(concept_art)
    im.save(f"{tmp_filename}.png")

  pdf = PDF()
  pdf.title = "Untitled" if title.strip() == "" else title
  pdf.art = None if concept_art is None else f"{tmp_filename}.png"
  pdf.print_chapter(editor.encode('latin-1', 'replace').decode('latin-1'))
  pdf.output(f'{tmp_filename}.pdf', 'F')

  return (
      gr.update(value=f'{tmp_filename}.pdf', visible=True),
      "   "
  )

def select(editor, evt: gr.SelectData):
  return [
      evt.value,
      evt.index[0],
      evt.index[1]
  ]

def get_gen_txt(title, editor, prompt, only_gen_text=False):
  if editor.strip() == '':
    ppm = get_new_ppm(f'Write a short paragraph to start a short story titled "{title}" for me')
  else:
    ppm = get_new_ppm(f"""{prompt}
--------------------------------
{editor}""")

  try:
    txt = gen_text(ppm.build_prompts(), hf_model=MODEL_ID, hf_token=TOKEN)
    if only_gen_text:
      return txt + "\n\n"
    else:
      return editor + txt + "\n\n"
  except ValueError as e:
    print(f"something went wrong - {e}")
    return editor

def gen_txt(title, editor, prompt):
  return [
      get_gen_txt(title, editor, "Write the next paragraph based on the following stories so far." if prompt.strip() == "" else prompt),
      0,
      gr.update(interactive=True),
      gr.update(visible=False),
      gr.update(visible=False),
      gr.update(visible=False),
      gr.update(interactive=True),
      gr.update(interactive=True),
  ]

def chat_gen(editor, chat_txt, chatbot, ppm, regen=False):
  ppm.ctx = f"""\
You are a helpful, respectful and honest assistant.

you must consider multi-turn conversations.

Answer to questions based on the written stories so far as below
----------------
{editor}
"""
  if regen:
    last_pingpong = ppm.pop_pingpong()
    chat_txt = last_pingpong.ping
  ppm.add_pingpong(PingPong(chat_txt, ''))

  try:
    txt = gen_text(ppm.build_prompts(), hf_model=MODEL_ID, hf_token=TOKEN)
    ppm.add_pong(txt)
  except ValueError as e:
    print(f"something went wrong - {e}")

  return [
    "",
    ppm.build_uis(),
    ppm
  ]

def chat(editor, chat_txt, chatbot, ppm):
  return chat_gen(editor, chat_txt, chatbot, ppm, regen=False) + [
      gr.update(interactive=True),
      gr.update(interactive=True)
  ]

def regen_chat(editor, chat_txt, chatbot, ppm):
  return chat_gen(editor, chat_txt, chatbot, ppm, regen=True)


def clear_chat():
  return [
      "",
      [],
      get_new_ppm_for_chat(),
      gr.update(interactive=False),
      gr.update(interactive=False)
  ]

def get_new_ppm_for_range():
  ppm = LLaMA2ChatPPManager()
  ppm.ctx = """\
  You are a helpful, respectful and honest writing helper. Always write text that suites to query.

  You DO NOT give explanation but just stories. DO NOT say such as 'Sure! Here's a short paragraph to start a short story:' or 'Sure, here is a revised version of ....:'
  """
  return ppm


def replace_sel(editor, replace_type, selected_text, sel_index_from, sel_index_to):
  ppm = get_new_ppm_for_range()

  ping = f"""replace {selected_text} in a single {replace_type} based on the story below
----------------
{editor}
"""

  ppm.add_pingpong(PingPong(ping, ''))

  try:
    txt = gen_text(ppm.build_prompts(), hf_model=MODEL_ID, hf_token=TOKEN)
    ppm.add_pong(txt)
  except ValueError as e:
    print(f"something went wrong - {e}")

  return [
      f"{editor[:sel_index_from]} {txt} {editor[sel_index_to:]}",
      "",
      0,
      0
  ]

def gen_alt(title, editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3):
  if num_enabled_alts < 3:
    gen_txt = get_gen_txt(title, editor, "Write the next paragraph based on the following stories so far.", only_gen_text=True)

  return [
      min(num_enabled_alts+1, 3),
      gr.update(interactive=False if num_enabled_alts >=2 else True),
      gr.update(visible=True if num_enabled_alts >=0 else False),
      gr.update(value=gen_txt if num_enabled_alts == 0 else alt_btn1),
      gr.update(visible=True if num_enabled_alts >=1 else False),
      gr.update(value=gen_txt if num_enabled_alts == 1 else alt_btn2),
      gr.update(visible=True if num_enabled_alts >=2 else False),
      gr.update(value=gen_txt if num_enabled_alts == 2 else alt_btn3),
      " ",
      gr.update(interactive=True),
      gr.update(interactive=True),
  ]

def fill_with_gen(alt_txt, editor):
  return [
      editor + alt_txt,
      0,
      gr.update(interactive=True),
      gr.update(visible=False),
      gr.update(visible=False),
      gr.update(visible=False)
  ]

def share(title, editor, cover_art):
  if title.strip() == '':
      title = 'Untitled'
    
  dataset = load_dataset(
      "chansung/llama2-stories", 
      download_mode=DownloadMode.FORCE_REDOWNLOAD,
      verification_mode=VerificationMode.NO_CHECKS
  )

  if cover_art is not None:
    buffered = io.BytesIO()
    cover_art = Image.fromarray(cover_art.astype('uint8'))    
    cover_art.save(buffered, format="JPEG")
    img_str = base64.b64encode(buffered.getvalue())
  else:
    img_str = ""
      
  dataset['train'] = dataset['train'].add_item({
      'title': title,
      'story': editor,
      'image': img_str
  })

  dataset.push_to_hub("chansung/llama2-stories", token=TOKEN)
    
  return  "       "

with gr.Blocks(css=STYLES) as demo:
  TITLE_PLACEHOLDERS = [
    'Pokemon training story',
    'The Sun of Shangri-La',
    'Man In The Future',
    'Friends',
    'Cyborg Of A Beast',
    'Man At The Graveyard',
    'Vampire Of The Land',
    'A software engineer who is looking for job',
    'A software engineer licensed to kill'
  ]
    
  num_enabled_alts = gr.State(0)

  sel_index_from = gr.State(0)
  sel_index_to = gr.State(0)

  chat_history = gr.State(get_new_ppm_for_chat())

  gr.Markdown("# Co-writing with AI", elem_classes=['center'])
  gr.Markdown(
      "This application is designed for you to collaborate with LLM to co-write stories. It is inspired by [Wordcraft project](https://wordcraft-writers-workshop.appspot.com/) from Google's PAIR and Magenta teams. "
      "This application built on [Gradio](https://www.gradio.app), and the underlying text generation is powered by [Hugging Face Inference API](https://huggingface.co/inference-api). The text generation model might"
      "be changed over time, but [meta-llama/Llama-2-70b-chat-hf](https://huggingface.co/meta-llama/Llama-2-70b-chat-hf) is selected for now. On the `exporting` tab, you can generate cover art for your story. You can "
      "design your own prompt to generate the cover art, or you can let the LLaMA2 to generate one for you. Currently, Stable Diffusion XL model `stabilityai/stable-diffusion-xl-base-1.0` by [Stability AI](https://stability.ai/) is "
      "used.",
      elem_classes=['center', 'small-big'])

  with gr.Row():
    with gr.Column(scale=5):
      editor = gr.Textbox(
          placeholder="Write your own story in this Textbox. Here are some guides for you. \n\n"
                      "1. If you don't know where to get started, go to the Chatting tab on the right pane. "
                      "You can discuss with AI about whatever you have in your mind. The first thing you would do "
                      "is settling down the Title of the story that you are going to write. \n\n"
                      "2. If you settle down the Title, click generate text on Control tab of the right pane. "
                      "Or just write the first short paragraph and click generate text button. Then, AI will suggest "
                      "the next paragraph for you. Whenever clicking Generate alternative button, it will suggest the "
                      "next paragraph to add. It will suggest up to 3 alternatives, so choose one that you like.\n\n"
                      "3. On the hatting tab on the right pane, you can discuss the written story. The whole text "
                      "from this textbox will be injected as a context. Understand and modify the story better with it. \n\n"
                      "4. On the Exporting/Sharing tab on the right pane, you can export your own stroy in PDF format. "
                      "Also, if you like sharing, you can share your work with community.",
          lines=32, max_lines=32, 
          elem_classes=['no-label', 'small-big-textarea']
      )
      word_counter = gr.Markdown("0 words", elem_classes=['right'])

    with gr.Column(scale=3):
      with gr.Tab("Control"):
        with gr.Column(elem_classes=['group-border']):
          gr.Markdown('### title')
          title = gr.Textbox(random.choice(TITLE_PLACEHOLDERS), elem_classes=['no-label'])
          random_title_btn = gr.Button("get random title", elem_classes=['control-label-font', 'control-button'])
        
        with gr.Column(elem_classes=['group-border']):
          with gr.Column():
            gr.Markdown("For instant generation and concatenation, use `generate text` button. "
                        "Want to explore alternative choices? use `generate alternatives` button.")

            with gr.Accordion("longer guideline", open=False):
              gr.Markdown("`generate text` button generate continued text and attach it to the end. "
                        "on the other hand, `generate alternatives` button generate alternate texts "
                        "up to 3 and let you choose one of them. In both cases, **Write the next paragraph based on "
                        "the following stories so far.** is the default prompt. If you want to try your own designed "
                        "prompt, enter it in the textbox below.")

            prompt = gr.Textbox(placeholder="design your own prompt", elem_classes=['no-label'])

          with gr.Row():
            gen_btn = gr.Button("generate text", elem_classes=['control-label-font', 'control-button'])
            gen_alt_btn = gr.Button("generate alternatives", elem_classes=['control-label-font', 'control-button'])

          with gr.Column():
            with gr.Row(visible=False) as first_alt:
              gr.Markdown("↳", scale=1, elem_classes=['wrap'])
              alt_btn1 = gr.Button("Alternative 1", elem_classes=['alt-button'], scale=8)

            with gr.Row(visible=False) as second_alt:
              gr.Markdown("↳", scale=1, elem_classes=['wrap'])
              alt_btn2 = gr.Button("Alternative 2", elem_classes=['alt-button'], scale=8)

            with gr.Row(visible=False) as third_alt:
              gr.Markdown("↳", scale=1, elem_classes=['wrap'])
              alt_btn3 = gr.Button("Alternative 3", elem_classes=['alt-button'], scale=8)

        with gr.Column(elem_classes=['group-border']):
          with gr.Row():
            selected_text = gr.Markdown("Selected text will be displayed in this area", elem_classes=['highlighted-text'])

          with gr.Row():
            with gr.Column(elem_classes=['no-gap']):
              replace_sel_btn = gr.Button("replace selection", elem_classes=['control-label-font', 'control-button'])
              replace_type = gr.Dropdown(choices=['word', 'sentense', 'phrase', 'paragraph'], value='sentense', interactive=True, elem_classes=['no-label'])

      with gr.Tab("Chatting"):
        with gr.Column(elem_classes=['group-border']):
          gr.Markdown("in this section, you can chat with LLaMA2 about the story that you are writing right now. "
                      "all the text in the textbox on the left will be injected as a context.")
          chatbot = gr.Chatbot([], elem_classes=['no-label-chatbot'])
          chat_txt = gr.Textbox(placeholder="enter question", elem_classes=['no-label'])

          with gr.Row():
            clear_btn = gr.Button("clear", interactive=False, elem_classes=['control-label-font', 'control-button'])
            regen_btn = gr.Button("regenerate", interactive=False, elem_classes=['control-label-font', 'control-button'])

      with gr.Tab("Exporting/Sharing"):
        with gr.Column(elem_classes=['group-border']):
          gr.Markdown("generate cover art with [`stabilityai/stable-diffusion-xl-base-1.0`](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) model. "
                      "design your own prompt in the textbox below, or just hit 'generate prompt for cover art` button. LLaMA2 "
                      "model will suggest a prompt for you based on your story.")

          cover_art = gr.Image(interactive=False, height=256, width=256, elem_classes=['no-label-image', 'margin-auto'])
          gen_cover_art_prompt = gr.Textbox(lines=3, max_lines=3, elem_classes=['no-label'])
          # toggle between "generate prompt for cover art" and "generate cover art"
          gen_cover_art_btn = gr.Button("generate prompt for cover art", elem_classes=['control-label-font', 'control-button'])

        with gr.Column(elem_classes=['group-border']):
          gen_pdf_btn = gr.Button("export as PDF", elem_classes=['control-label-font', 'control-button'])
          pdf_file = gr.File(visible=False)
            
          gr.Markdown("by clicking the button below, your stories will be shared with community. All shared stories will be "
                      "archived in [llama2-stories](https://huggingface.co/datasets/chansung/llama2-stories) dataset repository, and "
                      "they will be displayed in [chansung/LLaMA2-Story-Showcase](https://huggingface.co/spaces/chansung/LLaMA2-Story-Showcase) space.")
            
          share_btn = gr.Button("share", elem_classes=['control-label-font', 'control-button'])

  progress_bar = gr.Textbox(elem_classes=['no-label'])

  random_title_btn.click(
    fn=None,
    inputs=None,
    outputs=[title],
    _js="""
function get_random_title() {
  let titles = ['Pokemon training story', 'The Sun of Shangri-La', 'Man In The Future', 'Friends', 'Cyborg Of A Beast', 'Man At The Graveyard', 'Vampire Of The Land', 'A software engineer who is looking for job', 'A software engineer licensed to kill'];

  return titles[(Math.floor(Math.random() * titles.length))];
}
"""
  )
    
    
  share_btn.click(
      fn=share,
      inputs=[title, editor, cover_art],
      outputs=[progress_bar]
  )
    
  gen_pdf_btn.click(
      lambda t, e, c: generate_pdf(t, e, c),
      inputs=[title, editor, cover_art],
      outputs=[pdf_file, progress_bar]
  )

  gen_cover_art_btn.click(
      gen_art,
      inputs=[editor, cover_art, gen_cover_art_prompt],
      outputs=[cover_art, gen_cover_art_prompt]
  )

  gen_cover_art_prompt.change(
      fn=None,
      inputs=[gen_cover_art_prompt],
      outputs=[gen_cover_art_btn],
      _js="(t) => t.trim() == '' ? 'generate prompt for cover art' : 'generate cover art'"
  )

  editor.change(
      fn=None,
      inputs=[editor],
      outputs=[word_counter, selected_text],
      _js="(e) => [e.split(/\s+/).length, '']"
  )

  editor.select(
      fn=select,
      inputs=[editor],
      outputs=[selected_text, sel_index_from, sel_index_to],
      show_progress='minimal'
  )

  gen_btn.click(
      lambda: (
          gr.update(interactive=False),
          gr.update(interactive=False),
          gr.update(interactive=False),
          gr.update(interactive=False),
      ),
      inputs=None,
      outputs=[gen_btn, gen_alt_btn, replace_sel_btn]
  ).then(
      fn=gen_txt,
      inputs=[title, editor, prompt],
      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt, gen_btn, replace_sel_btn]
  )

  gen_alt_btn.click(
      lambda: (
          gr.update(interactive=False),
          gr.update(interactive=False),
          gr.update(interactive=False),
          gr.update(interactive=False),
      ),
      inputs=None,
      outputs=[gen_btn, gen_alt_btn, replace_sel_btn]
  ).then(
      fn=gen_alt,
      inputs=[title, editor, num_enabled_alts, alt_btn1, alt_btn2, alt_btn3],
      outputs=[num_enabled_alts, gen_alt_btn, first_alt, alt_btn1, second_alt, alt_btn2, third_alt, alt_btn3, progress_bar, gen_btn, replace_sel_btn],
  )

  alt_btn1.click(
      fn=fill_with_gen,
      inputs=[alt_btn1, editor],
      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt]
  )
  alt_btn2.click(
      fn=fill_with_gen,
      inputs=[alt_btn2, editor],
      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt]
  )
  alt_btn3.click(
      fn=fill_with_gen,
      inputs=[alt_btn3, editor],
      outputs=[editor, num_enabled_alts, gen_alt_btn, first_alt, second_alt, third_alt]
  )

  replace_sel_btn.click(
      fn=replace_sel,
      inputs=[editor, replace_type, selected_text, sel_index_from, sel_index_to],
      outputs=[editor, selected_text, sel_index_from, sel_index_to],
      show_progress='minimal'
  )

  chat_txt.submit(
      fn=chat,
      inputs=[editor, chat_txt, chatbot, chat_history],
      outputs=[chat_txt, chatbot, chat_history, clear_btn, regen_btn]
  )

  regen_btn.click(
      fn=regen_chat,
      inputs=[editor, chat_txt, chatbot, chat_history],
      outputs=[chat_txt, chatbot, chat_history]
  )

  clear_btn.click(
      fn=clear_chat,
      inputs=None,
      outputs=[chat_txt, chatbot, chat_history, clear_btn, regen_btn],
      show_progress='minimal'
  )

demo.queue(concurrency_count=5, max_size=256)
demo.launch()