import os
from threading import Thread
from typing import Iterator

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

DESCRIPTION = """\
# GRAG Llama-3.1-8B-SFT

👉 Quick test for multiple references.
"""

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

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

model_id = "avemio-digital/Llama3.1_8B_SFT_after_CPT_with_1_epoch"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    device_map="auto",
    torch_dtype=torch.bfloat16,
)
model.config.sliding_window = 4096
model.eval()
special_tokens = {'additional_special_tokens': ['<|im_start|>', '<|im_end|>']}
tokenizer.add_special_tokens(special_tokens)
model.resize_token_embeddings(len(tokenizer))
im_end_token_id = tokenizer.convert_tokens_to_ids('<|im_end|>')

system_message = """Beantworte die Fragen nur anhand des zur Verfügung gestellten Kontexts. Nachdem du die Frage beantwortet hast füge bitte den Index der Referenz aus dem Array 'context' in folgendem Format hinzu: [[relevant_index]]
Hier kommt ein Beispiel einer Anfrage und wie du darauf antworten sollst:
<Anfrage>
Hier würde die erste Frage stehen? Hier würde eine weitere Frage stehen?
Kontext:
'context': [{'id': 500017, 'source': 'Relevanter Kontext, der zur Beantwortung der ersten Frage genutzt wird'}, {'id': 500020, 'source': 'Kontext der nicht zur Beantwortung der Frage genutzt wird'}, {'id': 500018, 'source': 'Kontext der nicht zur Beantwortung der Frage genutzt wird'}, {'id': 500019, 'source': 'Relevanter Kontext, der zur Beantwortung der zweiten Frage genutzt wird'}]
</Anfrage>
<Deine Antwort>
Hier steht deine Antwort zur ersten Frage, welche nur Informationen aus dem relevanten Kontextabschnitt verwendet und in eine Antwort umformuliert.[[0]]
Anschließend steht hier deine Antwort zur zweiten Frage, welche nur Informationen aus dem für diese Frage relevanten Kontextabschnitt verwendet und in eine Antwort umformuliert.[[3]]
</Deine Antwort>
"""
@spaces.GPU(duration=90)
def generate(
    message: str,
    chat_history: list[dict],
    max_new_tokens: int = 2048,
    temperature: float = 0.6,
    top_p: float = 0.9,
    top_k: int = 50,
    repetition_penalty: float = 1.2,
) -> Iterator[str]:
    conversation = chat_history.copy()
    conversation.append({"role": "system", "content": system_message})
    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=20.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,
        num_beams=1,
        repetition_penalty=repetition_penalty,
        eos_token_id=im_end_token_id,
        pad_token_id=tokenizer.eos_token_id,
    )
    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,
    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.1,
            maximum=4.0,
            step=0.1,
            value=0.6,
        ),
        gr.Slider(
            label="Top-p (nucleus sampling)",
            minimum=0.05,
            maximum=1.0,
            step=0.05,
            value=0.9,
        ),
        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.2,
        ),
    ],
    stop_btn=None,
    examples=[
        ["""Wie lang ist die Talbrücke Rahmede und über welche Gewässer spannt sie sich? Wann musste die geplante Sprengung des Bauwerks aufgrund von Problemen beim Vergabeverfahren verschoben werden?
Kontext:
'context': [{'id': 1000034, 'source': 'Die Talbrücke Rahmede ist eine 453 Meter lange Brücke der Bundesautobahn 45. Sie liegt zwischen den Autobahnanschlüssen Lüdenscheid-Nord und Lüdenscheid. Das Bauwerk überspannt nördlich von Lüdenscheid das Tal der Rahmede, eines Zuflusses der Lenne, und die Landesstraße 530. Die Trasse der Autobahn verläuft im Bereich der Brücke im Grundriss gekrümmt.'}, {'id': 1000036, 'source': 'Die Talbrücke wurde von 1965 bis 1968 errichtet. Die Verkehrsplaner prognostizierten in den 1960er Jahren, dass täglich 25.000 Fahrzeuge sie nutzen würden. Ende der 2010er Jahre waren es 64.000 Fahrzeuge, davon 13.000 Lastkraftwagen. 1960 betrug das zulässige LKW-Gesamtgewicht 32 Tonnen und die Achslast 10 Tonnen, 1968 dann 38 Tonnen. Im Jahr 2021 waren maximal 40 Tonnen und 11,5 Tonnen Achslast, im kombinierten Verkehr bis 44 t Gesamtmasse zulässig. Hinzu kamen häufige illegale Überschreitungen von Gesamtgewicht und Achslasten.'}, {'id': 1000035, 'source': 'Das Bauwerk ist seit dem 2. Dezember 2021 wegen Schäden am Tragwerk dauerhaft gesperrt, eine Reparatur ist nicht möglich. Die Sprengung war zunächst für 2022 geplant, musste jedoch aufgrund von Problemen beim Vergabeverfahren verschoben werden.'}, {'id': 1000037, 'source': 'Wegen Mängeln muss die Brücke im Zuge des Ausbaus der Bundesautobahn 45 durch einen Neubau ersetzt werden. Das Bauwerk erhielt bei der Brückenhauptprüfung 2011 eine Zustandsnote 3,0 („nicht ausreichend“). Eine kurzfristige Instandsetzung sollte mit einem Zeitrahmen von einem Jahr bei Kosten von 18,67 Millionen Euro erfolgen. Die Maßnahme wurde aus Kostengründen mit Verweis auf den geplanten Neubau verworfen.'}]"""],
    ],
    cache_examples=False,
    type="messages",
)

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

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