from typing import Iterator

import gradio as gr
import torch

from model import run

DEFAULT_SYSTEM_PROMPT = ""
MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = 4000
DESCRIPTION = """
# Baichuan2-13B-Chat
Baichuan 2 is the new generation of open-source large language models launched by Baichuan Intelligent Technology. It was trained on a high-quality corpus with 2.6 trillion tokens.
"""
LICENSE = ""

if not torch.cuda.is_available():
  DESCRIPTION += '\n<p>Running on CPU 🥶 This demo does not work on CPU.</p>'


def clear_and_save_textbox(message: str) -> tuple[str, str]:
  return '', message

def display_input(
  message: str,
  history: list[tuple[str, str]]
) -> list[tuple[str, str]]:
  history.append((message, ''))
  return history

def delete_prev_fn(
  history: list[tuple[str, str]]
) -> tuple[list[tuple[str, str]], str]:
  try:
    message, _ = history.pop()
  except IndexError:
    message = ''
  return history, message or ''

def generate(
  message: str,
  history_with_input: list[tuple[str, str]],
  max_new_tokens: int,
  temperature: float,
  top_p: float,
  top_k: int,
) -> Iterator[list[tuple[str, str]]]:
  if max_new_tokens > MAX_MAX_NEW_TOKENS:
    raise ValueError

  history = history_with_input[:-1]
  generator = run(message, history, max_new_tokens, temperature, top_p, top_k)
  for response in generator:
    yield history + [(message, response)]

def process_example(message: str) -> tuple[str, list[tuple[str, str]]]:
  generator = generate(message, [], DEFAULT_MAX_NEW_TOKENS, 1, 0.95, 5)
  for x in generator:
    pass
  return '', x

with gr.Blocks(css='style.css') as demo:
  gr.Markdown(DESCRIPTION)
  gr.DuplicateButton(
    value='Duplicate Space for private use',
    elem_id='duplicate-button'
  )

  with gr.Group():
    chatbot = gr.Chatbot(label='Chatbot')
    with gr.Row():
      textbox = gr.Textbox(
        container=False,
        show_label=False,
        placeholder='Type a message...',
        scale=10,
      )
      submit_button = gr.Button(
        'Submit',
        variant='primary',
        scale=1,
        min_width=0
      )

  with gr.Row():
    retry_button = gr.Button('🔄  Retry', variant='secondary')
    undo_button = gr.Button('↩️ Undo', variant='secondary')
    clear_button = gr.Button('🗑️  Clear', variant='secondary')

  saved_input = gr.State()

  with gr.Accordion(label='Advanced options', open=False):
    max_new_tokens = gr.Slider(
      label='Max new tokens',
      minimum=1,
      maximum=MAX_MAX_NEW_TOKENS,
      step=1,
      value=DEFAULT_MAX_NEW_TOKENS,
    )
    temperature = gr.Slider(
      label='Temperature',
      minimum=0.1,
      maximum=4.0,
      step=0.1,
      value=1.0,
    )
    top_p = gr.Slider(
      label='Top-p (nucleus sampling)',
      minimum=0.05,
      maximum=1.0,
      step=0.05,
      value=0.95,
    )
    top_k = gr.Slider(
      label='Top-k',
      minimum=1,
      maximum=1000,
      step=1,
      value=5,
    )

  gr.Examples(
    examples=[
      '介绍下你自己',
      '找到下列数组的中位数[3.1,6.2,1.3,8.4,10.5,11.6,2.1]，请用python代码完成以上功能',
      '鸡和兔在一个笼子里，共有26个头，68只脚，那么鸡有多少只，兔有多少只？',
      '以下物理常识题目，哪一个是错误的?A.在自然环境下，声音在固体中传播速度最快。B.牛顿第一定律:一个物体如果不受力作用，将保持静止或匀速直线运动的状态。C.牛顿第三定律:对于每个作用力，都有一个相等而反向的反作用力。D.声音在空气中的传播速度为1000m/s。',
    ],
    inputs=textbox,
    outputs=[textbox, chatbot],
    fn=process_example,
    cache_examples=True,
  )

  gr.Markdown(LICENSE)

  textbox.submit(
    fn=clear_and_save_textbox,
    inputs=textbox,
    outputs=[textbox, saved_input],
    api_name=False,
    queue=False,
  ).then(
    fn=display_input,
    inputs=[saved_input, chatbot],
    outputs=chatbot,
    api_name=False,
    queue=False,
  ).then(
    fn=generate,
    inputs=[
      saved_input,
      chatbot,
      max_new_tokens,
      temperature,
      top_p,
      top_k,
    ],
    outputs=chatbot,
    api_name=False,
  )

  button_event_preprocess = submit_button.click(
    fn=clear_and_save_textbox,
    inputs=textbox,
    outputs=[textbox, saved_input],
    api_name=False,
    queue=False,
  ).then(
    fn=display_input,
    inputs=[saved_input, chatbot],
    outputs=chatbot,
    api_name=False,
    queue=False,
  ).then(
    fn=generate,
    inputs=[
      saved_input,
      chatbot,
      max_new_tokens,
      temperature,
      top_p,
      top_k,
    ],
    outputs=chatbot,
    api_name=False,
  )

  retry_button.click(
    fn=delete_prev_fn,
    inputs=chatbot,
    outputs=[chatbot, saved_input],
    api_name=False,
    queue=False,
  ).then(
    fn=display_input,
    inputs=[saved_input, chatbot],
    outputs=chatbot,
    api_name=False,
    queue=False,
  ).then(
    fn=generate,
    inputs=[
      saved_input,
      chatbot,
      max_new_tokens,
      temperature,
      top_p,
      top_k,
    ],
    outputs=chatbot,
    api_name=False,
  )

  undo_button.click(
    fn=delete_prev_fn,
    inputs=chatbot,
    outputs=[chatbot, saved_input],
    api_name=False,
    queue=False,
  ).then(
    fn=lambda x: x,
    inputs=[saved_input],
    outputs=textbox,
    api_name=False,
    queue=False,
  )

  clear_button.click(
    fn=lambda: ([], ''),
    outputs=[chatbot, saved_input],
    queue=False,
    api_name=False,
  )

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