Change SDK

README.md

@@ -3,8 +3,8 @@ title: Bloom Chat
 emoji: ⚡
 colorFrom: purple
 colorTo: green
-sdk: gradio
-sdk_version: 3.4
+sdk: streamlit
+sdk_version: 1.10.0
 app_file: app.py
 pinned: false
 license: openrail

app.py

@@ -1,27 +1,104 @@
-# import streamlit as st
+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
 
-# text = st.text_area("Prefix", value="DM: You enter the room.")
-# batch = st.number_input("Variants", value=5)
-# st.markdown(f"{text} {batch}")
 
+@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
 
-import streamlit as st
-from transformers import pipeline
-from PIL import Image
+    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)
 
-pipeline = pipeline(task="image-classification", model="julien-c/hotdog-not-hotdog")
+t = st.empty()
+firstline = ""
 
-st.title("Hot Dog? Or Not?")
+def PrintSome(text):
+    global t, firstline
+    firstline += text
+    t.markdown(f"{firstline}...")
 
-file_name = st.file_uploader("Upload a hot dog candidate image")
+Sureprint("before inference")
+choices, batch_time = PrintContinuation(text, model, PrintSome, batch, 50)
+Sureprint("after inference")
 
-if file_name is not None:
-  col1, col2 = st.columns(2)
+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}"
 
-  image = Image.open(file_name)
-  col1.image(image, use_column_width=True)
-  predictions = pipeline(image)
+t.markdown(final_page)
 
-  col2.header("Probabilities")
-  for p in predictions:
-    col2.subheader(f"{ p['label'] }: { round(p['score'] * 100, 1)}%")
+Sureprint("all done")