import torch
import transformers
import numpy as np
from huggingface_hub import hf_hub_download


tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")

hf_hub_download("OpenDungeon/gpt-j-8bit-ffbgem", "model.pt")

qmodel = torch.load("model.pt")

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

    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:
                print()
                print((time.perf_counter() - start_time) / (i - 4), "s per token")
                break

            input_dict = dict(input_ids=token_ix)
    print()
    return output
    
import streamlit as st

def process(text):
    return text[::-1]


text = st.text_area("Prompt")
t.markdown(f"## {process(text)[0:i]}...")