Spaces:
Runtime error
Runtime error
File size: 3,404 Bytes
07d1ab1 047c577 07d1ab1 047c577 07d1ab1 047c577 07d1ab1 047c577 07d1ab1 047c577 07d1ab1 047c577 07d1ab1 047c577 07d1ab1 09c0a40 07d1ab1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 |
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")
|