import torch
import transformers
import time
from huggingface_hub import snapshot_download
import streamlit as st
import copy
from transformers import AutoConfig, GPTJForCausalLM
from transformers.models.gptj.modeling_gptj import GPTJBlock
from tqdm import trange


@st.cache(allow_output_mutation=True)
def load_model():
    for down in trange(1, disable=True):
        fpath = snapshot_download("OpenDungeon/gpt-j-8bit-ffbgem", revision="separate")
    config = AutoConfig.from_pretrained("EleutherAI/gpt-j-6B")
    qconfig = torch.quantization.get_default_qconfig('fbgemm')
    torch.backends.quantized.engine = 'fbgemm'
    n_layer, config.n_layer = config.n_layer, 0

    model = GPTJForCausalLM(config)
    model.load_state_dict(torch.load(fpath + "/blocks/base.pt"))
    ref_block = torch.quantization.quantize_dynamic(
            GPTJBlock(config),
            {torch.nn.Linear: qconfig},
            dtype=torch.qint8,
            inplace=True
        )

    for i in trange(n_layer):
        new_block = copy.deepcopy(ref_block)
        new_block.load_state_dict(torch.load(f"{fpath}/blocks/block{i}.pt"))
        model.transformer.h.append(new_block)

    config.n_layer = len(model.transformer.h)
    del ref_block

    return transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"), model


def PrintContinuation(prompt, local_model, single_hook=None, batch=1, limit_tokens = 50):
    past_key_values = None  # used to keep track of conversation history
    input_dict = tokenizer([prompt] * batch, return_tensors='pt', padding=False)
    output = [""] * batch
    batch_time = 0
    
    with torch.inference_mode():
        for i in range(limit_tokens + 20):
            if i == 5:
                start_time = time.perf_counter()

            outputs = local_model.forward(**input_dict, use_cache=True, past_key_values=past_key_values)
            last_logits = outputs.logits[:, -1]

            for j in range(batch):
                last_logits[j, last_logits[j].topk(k=10).indices] += 10

            past_key_values = outputs.past_key_values
            token_ix = torch.multinomial(last_logits.softmax(-1), 1)
            output = [stream + tokenizer.decode(ix) for stream, ix in zip(output, token_ix)]

            if single_hook is not None:
                single_hook(tokenizer.decode(token_ix[0]))
            if i == limit_tokens:
                batch_time = (time.perf_counter() - start_time) / (i - 4)
                break

            input_dict = dict(input_ids=token_ix)
    return output, batch_time

import sys

def Sureprint(text):
    text = f"\nDDBG: {text}\n"
    print(text, flush=True)
    print(text, file=sys.stderr, flush=True)

Sureprint("ready to load")
tokenizer, model = load_model()
Sureprint("loaded")
text = st.text_area("Prefix", value="DM: You enter the room.")
Sureprint(f"text acquired '{text}'")
batch = st.number_input("Variants", value=5)

t = st.empty()
firstline = ""

def PrintSome(text):
    global t, firstline
    firstline += text
    t.markdown(f"{firstline}...")

Sureprint("before inference")
choices, batch_time = PrintContinuation(text, model, PrintSome, batch, 50)
Sureprint("after inference")

final_page = ""
for i in range(batch):
    final_page += f"#### choice №{i + 1}  \n{choices[i]}  \n______  \n"
final_page += f"Seconds per batch: {batch_time}, Batch: {batch}"

t.markdown(final_page)

Sureprint("all done")