Update app.py

app.py

@@ -1,48 +1,51 @@
-from transformers import BloomTokenizerFast, BloomModel
 import torch
-import gradio as gr
-
-mname = "bigscience/bloom-1b7"
-tokenizer = BloomTokenizerFast.from_pretrained(mname, use_cache=True)
-model = BloomModel.from_pretrained(mname, use_cache=True)
-
-def take_last_tokens(inputs, note_history, history):
-    """Filter the last 256 tokens"""
-    if inputs['input_ids'].shape[1] > 256:
-        inputs['input_ids'] = torch.tensor([inputs['input_ids'][0][-256:].tolist()])
-        inputs['attention_mask'] = torch.tensor([inputs['attention_mask'][0][-256:].tolist()])
-        note_history = ['</s> <s>'.join(note_history[0].split('</s> <s>')[2:])]
-        history = history[1:]
-    return inputs, note_history, history
-
-def add_note_to_history(note, note_history):
-    """Add a note to the historical information"""
-    note_history.append(note)
-    note_history = '</s> <s>'.join(note_history)
-    return [note_history]
-
-def chat(message, history):
-    history = history or []
-    if history: 
-        history_useful = ['</s> <s>'.join([str(a[0])+'</s> <s>'+str(a[1]) for a in history])]
-    else:
-        history_useful = []
-    history_useful = add_note_to_history(message, history_useful)
-    inputs = tokenizer(history_useful, return_tensors="pt")
-    inputs, history_useful, history = take_last_tokens(inputs, history_useful, history)
-    reply_ids = model.generate(**inputs)
-    response = tokenizer.batch_decode(reply_ids, skip_special_tokens=True)[0]
-    history_useful = add_note_to_history(response, history_useful)
-    list_history = history_useful[0].split('</s> <s>')
-    history.append((list_history[-2], list_history[-1]))
-    return history, history
-
-gr.Interface(
-    fn=chat,
-    theme="huggingface",
-    css=".footer {display:none !important}",
-    inputs=["text", "state"],
-    outputs=["message", "state"],
-    title="Bloom 1b3 chat",
-    allow_flagging="never",
-    ).launch()
+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]}...")