#!/usr/bin/env python

import os
from threading import Thread
from typing import Iterator
import torch

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

MAX_MAX_NEW_TOKENS = 2048
DEFAULT_MAX_NEW_TOKENS = 1024
MAX_INPUT_TOKEN_LENGTH = int(os.getenv("MAX_INPUT_TOKEN_LENGTH", "8192"))

if torch.cuda.is_available():
    model_id = "BramVanroy/GEITje-ultra-dpo-5e-7lr-128tbs-0.1b"
    model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.bfloat16, device_map="auto")
    tokenizer = AutoTokenizer.from_pretrained(model_id)


# @spaces.GPU
def generate(
    message: str,
    chat_history: list[tuple[str, str]],
    max_new_tokens: int = 1024,
    temperature: float = 1,
    top_p: float = 1.,
    top_k: int = 50,
    do_sample: bool = False,
    repetition_penalty: float = 1.,
) -> Iterator[str]:
    conversation = [{"role": "system", "content": ""}]
    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")
    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.")
    input_ids = input_ids.to(model.device)

    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=do_sample,
        top_p=top_p,
        top_k=top_k,
        temperature=temperature,
        num_beams=1,
        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)


chat_interface = gr.ChatInterface(
    fn=generate,
    chatbot=gr.Chatbot(height=450,
                      label="GEITje-ultra",
                      show_share_button=True,
                      avatar_images=(None, 'geitje-ultra-avatar.png')),
    additional_inputs=[
        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.05,
            maximum=2,
            step=0.05,
            value=1.0,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=1.0,
        ),
        gr.Slider(
            label="Top-k",
            minimum=1,
            maximum=1000,
            step=1,
            value=50,
        ),
        gr.Slider(
            label="Repetition penalty",
            minimum=1.0,
            maximum=2.0,
            step=0.05,
            value=1.,
        ),
        gr.Checkbox(
            label="Do sample",
            value=False,
        )
    ],
    examples=[
        ["""Vraagje: welk woord hoort er niet in dit rijtje thuis: "auto, vliegtuig, geit, bus"?"""],
        ["Schrijf een nieuwsbericht voor De Speld over de inzet van een kudde geiten door het Nederlands Forensisch Instituut"],
        ["Wat zijn 3 leuke dingen om te doen als ik een weekendje naar Limburg ga?"],
        ["Met wie trad clown Bassie op?"],
        ["Kan je naar de maan fietsen?"],
        ["Wat is het belang van open-source taalmodellen?"],
        ["""```
# Wortelverkopers krijgen miljoenenboete voor ongeoorloofd samenspannen

Toezichthouder ACM heeft een Nederlands wortelkartel aangepakt. Vier telers en verkopers van wortelen krijgen samen ruim 2,5 miljoen euro boete vanwege ongeoorloofde afspraken over het verdelen van de markt.
Het gaat om telers en verkopers Laarakker, VanRijsingen, Veco en Verduyn. De vier bedrijven verkopen waspeen en Parijse wortelen aan conserven- en diepvriesfabrikanten in Nederland, België en Duitsland. Waspeen wordt vaak verkocht in potten of blikken in een mix met erwtjes.

De vier bedrijven hadden in 2018 afgesproken dat ze tien jaar lang niet overal de concurrentie met elkaar zouden aangaan. Zo zou Veco tien jaar lang geen waspeen telen of verkopen. Daarnaast zouden Laarakker, VanRijsingen en Verduyn juist de Parijse wortelen links laten liggen.

Ook betaalden de andere wortelverkopers Veco ter compensatie van de afspraken. Laarakker en Veco maakten ook nog afzonderlijke afspraken over de levering van Parijse wortelen aan Duitse klanten.

Zulke afspraken zijn verboden. Als concurrentie door die samenwerking achterwege blijft en er dus sprake is van een kartel, betalen kopers mogelijk een hogere prijs, stelt de ACM.

Twee van de wortelbedrijven werkten mee door meer informatie over de ongeoorloofde afspraken te delen met de toezichthouder. Daardoor kregen zij een lagere boete.
```
Vat bovenstaand artikel kort samen."""]
    ],
    title="🐐 GEITje ultra 🤖",
    description="""\
<a href="https://huggingface.co/BramVanroy/GEITje-ultra">GEITje ultra</a> is een geavanceerde versie van GEITje, verder getraind op uitgebreide chat datasets en ook op preferentiedatasets om beter te aligneren met het gedrag van een gewenste chatbot, in dit geval gpt-4-turbo. Ik heb gemerkt dat dit ertoe leidt dat GEITje ultra veel uitgebreidere antwoorden geeft dan de originele GEITje chat. Dat kan positief of negatief zijn, afhankelijk van wat je verwachtingen zijn. Ik heb ook gemerkt dat GEITje ultra na verloop van tijd slechter wordt in z'n Nederlands in een gesprek. Dat bevestigt de noozaak aan een model dat van het begin gepretraind werd op het Nederlands!

Het is goed mogelijk dat het model alsnog hallucineert met nonsens-karakters omdat het DPO trainen té hard z'n werk gedaan heeft. De huidige versie lijkt robuuster dan vorige pogingen. Als je problemen ondervindt, [rapporteer die dan gerust](https://huggingface.co/spaces/BramVanroy/GEITje-ultra/discussions)! Dit blijt een onderzoeks+hobby-project - verwacht er dus geen production-ready model van.

Hoewel dit model gealigneerd is met AI feedback (gpt-4-turbo), kan en zal het nog steeds fouten maken, leugens vertellen, redeneringsfouten maken, en misschien wel stoute dingen vertellen. Gebruik dit model dus op eigen risico en controleer de output zelf!

<small>Deze demo is  sterk gebaseerd op de originele <a href="https://huggingface.co/spaces/Rijgersberg/GEITje-7B-chat">Geitje Chat</a> demo, die dan weer op deze [mistral demo](https://huggingface.co/spaces/hysts/mistral-7b) gebaseerd is.</small>
""",
    submit_btn="Genereer",
    stop_btn="Stop",
    retry_btn="🔄 Opnieuw",
    undo_btn="↩️ Ongedaan maken",
    clear_btn="🗑️ Wissen",
)

with gr.Blocks(css="style.css") as demo:
    chat_interface.render()

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