# pylint: skip-file

import subprocess

subprocess.run(
    f"pip install flash-attn --no-build-isolation",
    env={"FLASH_ATTENTION_SKIP_CUDA_BUILD": "TRUE"},
    shell=True,
)

import os
from threading import Thread
from typing import Iterator

import gradio as gr
import spaces
import torch
import SelfExtend
from transformers import AutoModelForCausalLM, AutoTokenizer, TextIteratorStreamer


MAX_MAX_NEW_TOKENS = 4096
DEFAULT_MAX_NEW_TOKENS = 1536
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "123392"))

DESCRIPTION = """\
# Playground with Ghost 8B Beta (p)

**Ghost 8B Beta** is a large language model developed with goals that include excellent multilingual support, superior knowledge capabilities, and cost-effectiveness. The model comes in two context length versions, 8k and 128k, along with multilingual function tools support by default. 

The languages supported are 🇺🇸 English, 🇫🇷 French, 🇮🇹 Italian, 🇪🇸 Spanish, 🇵🇹 Portuguese, 🇩🇪 German, 🇻🇳 Vietnamese, 🇰🇷 Korean and 🇨🇳 Chinese.

📋 Note: current model version is "disl-0x5" (10 Jul 2024), context length 128k (123392 tokens) and current status is "moderating / previewing". For detailed information about the model, see [here](https://ghost-x.org/docs/models/ghost-8b-beta/). Try to experience it the way you want!
"""


PLACEHOLDER = """
<div style="padding: 30px; text-align: center; display: flex; flex-direction: column; align-items: center;">
   <h1 style="font-size: 26px; margin-bottom: 2px; opacity: 0.20;">👻 Ghost 8B Beta</h1>
   <p style="font-size: 18px; margin-bottom: 2px; opacity: 0.10;">Ask and share whatever you want ~</p>
</div>
"""

LICENSE = """
<p/>

---
Ghost 8B Beta may give inaccurate information, including information about people, so please verify Ghost 8B Beta's answers. [Ghost 8B Beta](https://ghost-x.org/docs/models/ghost-8b-beta/) by [Ghost X](https://ghost-x.org).
"""

EXAMPLES = [
    [
        "What is the significance of the Higgs boson in the Standard Model of particle physics?"
    ],
    [
        "Qu'est-ce que l'effet fondateur et comment influence-t-il la diversité génétique d'une population?"
    ],
    ["Qual è il principio di Le Chatelier e come si applica agli equilibri chimici?"],
    [
        "¿Qué es una supernova y cuál es su importancia en la formación de elementos pesados en el universo?"
    ],
    [
        "Qual é a definição formal de uma integral de linha e como é utilizada em física?"
    ],
    [
        "Was versteht man unter dem Moho-Diskontinuität und welche Bedeutung hat sie für das Verständnis der Erdkruste?"
    ],
    [
        "Hiện tượng nhà kính là gì và nó ảnh hưởng như thế nào đến biến đổi khí hậu toàn cầu?"
    ],
    [
        "알고리즘의 시간 복잡도가 중요한 이유는 무엇이며, 시간 복잡도를 어떻게 분석하나요?"
    ],
    ["什么是CRISPR-Cas9基因编辑技术，它在现代生物学研究中的作用是什么？"],
    [
        "Create a Python function that takes a list of integers and returns the list sorted in ascending order without using the built-in sort or sorted functions."
    ],
    [
        "Écrivez une fonction en C++ qui trouve le plus long sous-tableau contigu avec une somme égale à zéro."
    ],
    [
        "Scrivi una funzione in Java che calcola il fattoriale di un numero utilizzando la ricorsione."
    ],
    [
        "Desarrolla una función en JavaScript que determine si una cadena de texto es un palíndromo, ignorando espacios y signos de puntuación."
    ],
    ["Implemente uma função em C# que verifique se uma matriz quadrada é simétrica."],
    [
        "Schreiben Sie eine Funktion in Swift, die eine gegebene Zeichenfolge in umgekehrter Reihenfolge zurückgibt, ohne integrierte Funktionen zu verwenden."
    ],
    [
        "Viết một hàm trong PHP để tìm tất cả các số nguyên tố trong một khoảng cho trước."
    ],
    [
        "파이썬을 사용하여 주어진 이진 트리가 이진 탐색 트리인지 확인하는 함수를 작성하십시오."
    ],
    [
        "用 Go 语言编写一个函数，计算给定字符串中每个字符出现的次数，并返回一个包含字符及其出现次数的映射。"
    ],
    [
        "Can you help me design a detailed project plan for developing a machine learning model for predicting stock prices?"
    ],
    [
        "Pouvez-vous m'aider à organiser un emploi du temps hebdomadaire pour maximiser la productivité de mon équipe de développement logiciel?"
    ],
    [
        "Puoi aiutarmi a creare un piano di sviluppo per un'applicazione mobile che gestisce le prenotazioni di ristoranti?"
    ],
    [
        "¿Podrías ayudarme a elaborar un plan detallado para la implementación de un sistema de gestión de contenido (CMS) en una empresa mediana?"
    ],
    [
        "Você pode me ajudar a planejar uma estratégia de desenvolvimento para um sistema de comércio eletrônico escalável?"
    ],
    [
        "Können Sie mir helfen, einen detaillierten Zeitplan für die Implementierung eines neuen ERP-Systems in unserem Unternehmen zu erstellen?"
    ],
    [
        "Bạn có thể giúp tôi xây dựng một kế hoạch phát triển chi tiết cho dự án xây dựng hệ thống quản lý chuỗi cung ứng không?"
    ],
    [
        "신경망 기반 이미지 인식 모델 개발을 위한 세부 프로젝트 계획을 세우는 데 도움을 줄 수 있나요?"
    ],
    ["你能帮我制定一个详细的开发计划，用于创建一个基于区块链的分布式账本系统吗？"],
    [
        "Prove that the sum of the squares of any two sides of a right triangle is equal to the square of the hypotenuse."
    ],
    [
        "Calculez la force gravitationnelle entre deux masses de 10 kg chacune séparées par une distance de 1 mètre."
    ],
    [
        "Determina la formula molecolare di un composto che contiene il 40% di carbonio, il 6.67% di idrogeno e il 53.33% di ossigeno in massa."
    ],
    [
        "Explica la teoría del ciclo económico de Schumpeter y cómo se aplica a la economía moderna."
    ],
    [
        "Calcule a energia potencial gravitacional de um objeto de 5 kg a uma altura de 10 metros acima do solo (g = 9,8 m/s²)."
    ],
    [
        "Beweisen Sie, dass jede Primzahl der Form 4k+1 als Summe zweier Quadrate geschrieben werden kann."
    ],
    [
        "Tính nồng độ mol của dung dịch H₂SO₄ khi hoà tan 98 gam H₂SO₄ vào nước để được 1 lít dung dịch."
    ],
    ["케인스 경제학의 핵심 개념과 그것이 현대 경제 정책에 미치는 영향을 설명하십시오."],
    ["计算一个质量为2 kg的物体在3米高处的重力势能（g = 9.8 m/s²）。"],
    [
        'Identify the author of a novel that features a dystopian society where "Big Brother" watches over its citizens and the protagonist works for the Ministry of Truth.'
    ],
    [
        "Quel est le seul mammifère capable de voler activement, souvent associé à la nuit et capable d'écholocalisation?"
    ],
    [
        "Qual è l'opera letteraria italiana che narra il viaggio immaginario di un poeta attraverso Inferno, Purgatorio e Paradiso, guidato da Virgilio e Beatrice?"
    ],
    [
        "¿Qué insecto es conocido por su organización social compleja, su capacidad para producir miel y su comunicación mediante la danza?"
    ],
    [
        "Qual é o fenômeno atmosférico que ocorre quando uma massa de ar quente se encontra com uma massa de ar frio, resultando em uma violenta tempestade giratória?"
    ],
    [
        "Welches literarische Werk beschreibt die Geschichte eines jungen Mädchens, das durch einen Kaninchenbau in eine fantastische Welt voller skurriler Charaktere fällt?"
    ],
    [
        "Động vật nào có thể tái sinh toàn bộ cơ thể từ một mảnh nhỏ của chính nó, thường sống dưới nước và có thể có nhiều xúc tu?"
    ],
    [
        "어떤 자연 현상은 태양빛이 대기 중의 물방울에 반사되고 굴절되어 발생하며, 하늘에 나타나는 여러 색깔의 아치 형태를 띠나요?"
    ],
    ["这部文学作品讲述了一位绅士和他的侍从的冒险故事，他们在"],
    [
        "Can you derive the Euler-Lagrange equation from the principle of stationary action in classical mechanics?"
    ],
    [
        "Expliquez la notion de « différence ontologique » chez Martin Heidegger et son importance pour la phénoménologie."
    ],
    [
        "Qual è il significato simbolico del colore blu nei dipinti di Giotto di Bondone durante il Rinascimento?"
    ],
    [
        "¿Cómo afecta el cambio de código a la estructura gramatical en comunidades bilingües de habla español-inglés?"
    ],
    [
        "Qual é o impacto da política monetária não convencional no controle da inflação durante uma crise econômica?"
    ],
    [
        "Erklären Sie den Unterschied zwischen deterministischen und nicht-deterministischen endlichen Automaten und ihre Anwendungsbereiche."
    ],
    [
        "Giải thích cơ chế của quá trình phiên mã ngược (reverse transcription) và tầm quan trọng của nó trong nghiên cứu HIV/AIDS."
    ],
    ["조선시대 성리학이 한국 사회와 문화에 미친 영향을 설명하세요."],
    ["如何解释量子纠缠现象，以及它在量子计算中的潜在应用？"],
    [
        "How can you design a daily schedule that maximizes productivity for a remote worker who has multiple meetings and project deadlines?"
    ],
    [
        "Quels sont les meilleures stratégies pour gérer les conflits au sein d'une équipe multiculturelle travaillant sur un projet commun?"
    ],
    [
        "Quali sono i migliori consigli per mantenere un equilibrio tra vita professionale e vita privata in un ambiente lavorativo stressante?"
    ],
    [
        "¿Cómo se puede elaborar un plan financiero personal efectivo que incluya ahorro para la jubilación, inversión y manejo de deudas?"
    ],
    [
        "Quais são as melhores práticas para implementar metodologias ágeis em uma equipe de desenvolvimento de software?"
    ],
    [
        "Welche Strategien können verwendet werden, um ein starkes berufliches Netzwerk aufzubauen und zu pflegen, insbesondere in der Tech-Branche?"
    ],
    [
        "Những bước nào cần thiết để xây dựng một lộ trình phát triển sự nghiệp bền vững trong lĩnh vực công nghệ thông tin?"
    ],
    ["프로젝트의 범위 변동을 효과적으로 관리하기 위한 최고의 방법은 무엇인가요?"],
    ["在快速变化的职场环境中，如何有效地实现工作与生活的平衡？"],
    [
        "Write an argumentative essay discussing the pros and cons of artificial intelligence in the workplace, including potential ethical concerns."
    ],
    [
        "Analysez les impacts sociaux et économiques de la digitalisation sur les petites entreprises en France."
    ],
    [
        "Scrivi un'email formale al direttore di una rivista per proporre un articolo sulla sostenibilità ambientale nelle città italiane."
    ],
    [
        "Elabora un informe detallado sobre los efectos del cambio climático en la biodiversidad de la región amazónica."
    ],
    [
        "Analise criticamente os principais pontos abordados no relatório anual do Banco Mundial sobre a pobreza global."
    ],
    [
        "Erstellen Sie eine technische Dokumentation für die Implementierung eines neuen Software-Features in einer bestehenden Anwendung."
    ],
    [
        "Viết một bài luận phân tích về tác động của cuộc cách mạng công nghiệp 4.0 đối với thị trường lao động Việt Nam."
    ],
    [
        "인공지능의 윤리적 문제에 대한 연구 논문을 작성하고, 다양한 사례를 통해 그 영향을 분석하세요."
    ],
    ["分析鲁迅的小说《阿Q正传》中反映的中国社会问题和作者的批判态度。"],
]

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


if torch.cuda.is_available():
    model_id = "lamhieu/ghost-8b-beta-disl-0x5-8k"
    model_tk = os.getenv("HF_TOKEN", None)
    model = AutoModelForCausalLM.from_pretrained(
        model_id,
        device_map="auto",
        torch_dtype=torch.bfloat16,
        attn_implementation="flash_attention_2",
        trust_remote_code=True,
        token=model_tk,
    )
    tokenizer = AutoTokenizer.from_pretrained(
        model_id,
        trust_remote_code=True,
        token=model_tk,
    )
    SelfExtend.apply(
        model,
        group_size=16,
        window_size=512,
        enable_flash_attention=True,
        flash_attention_impl="flash_attn",
    )
    model.generation_config.max_length = 123392


@spaces.GPU(duration=120)
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    system_prompt: str,
    max_new_tokens: int = 1536,
    temperature: float = 0.4,
    top_p: float = 0.95,
    top_k: int = 50,
    repetition_penalty: float = 1.0,
) -> Iterator[str]:
    conversation = []
    if system_prompt:
        conversation.append({"role": "system", "content": system_prompt})
    for user, assistant in chat_history:
        conversation.extend(
            [
                {"role": "user", "content": user},
                {"role": "assistant", "content": assistant},
            ]
        )
    conversation.append({"role": "user", "content": message})

    input_ids = tokenizer.apply_chat_template(
        conversation, add_generation_prompt=True, return_tensors="pt"
    )
    input_ids = input_ids.to(model.device)
    if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
        input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
        gr.Warning(
            f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens."
        )

    streamer = TextIteratorStreamer(
        tokenizer, timeout=10.0, skip_prompt=True, skip_special_tokens=True
    )
    generate_kwargs = dict(
        input_ids=input_ids,
        streamer=streamer,
        max_new_tokens=max_new_tokens,
        do_sample=True,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        repetition_penalty=repetition_penalty,
    )
    t = Thread(target=model.generate, kwargs=generate_kwargs)
    t.start()

    outputs = []
    for text in streamer:
        outputs.append(text)
        yield "".join(outputs)


chatbot = gr.Chatbot(height=500, placeholder=PLACEHOLDER, label="Ghost 8B Beta")

chat_interface = gr.ChatInterface(
    fn=generate,
    chatbot=chatbot,
    fill_height=True,
    additional_inputs=[
        gr.Textbox(label="System prompt", lines=6),
        gr.Slider(
            label="Max new tokens",
            minimum=1,
            maximum=MAX_MAX_NEW_TOKENS,
            step=1,
            value=DEFAULT_MAX_NEW_TOKENS,
        ),
        gr.Slider(
            label="Temperature",
            minimum=0.1,
            maximum=2.0,
            step=0.1,
            value=0.4,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.95,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=100,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.0,
        ),
    ],
    stop_btn="Stop",
    cache_examples=False,
    examples=EXAMPLES,
    examples_per_page=9,
)

with gr.Blocks(fill_height=True, css="style.css") as demo:
    gr.Markdown(DESCRIPTION)
    chat_interface.render()
    gr.Markdown(LICENSE)

if __name__ == "__main__":
    demo.queue(max_size=20).launch(share=True)
    # demo.launch(share=True)