Cache + Measure time

app.py

@@ -1,20 +1,23 @@
 import torch
 import transformers
-import numpy as np
+import time
 from huggingface_hub import hf_hub_download
+import streamlit as st
 
 
-tokenizer = transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B")
-
-hf_hub_download("OpenDungeon/gpt-j-8bit-ffbgem", "model.pt")
+@st.cache
+def load_model():
+    hf_hub_download("OpenDungeon/gpt-j-8bit-ffbgem", "model.pt")
+    qmodel = torch.load("model.pt")
+    return transformers.AutoTokenizer.from_pretrained("EleutherAI/gpt-j-6B"), qmodel
 
-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
-
+    batch_time = 0
+    
     with torch.inference_mode():
         for i in range(limit_tokens + 20):
             if i == 5:
@@ -33,16 +36,25 @@ def PrintContinuation(prompt, local_model, single_hook=None, batch=1, limit_toke
             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")
+                batch_time = (time.perf_counter() - start_time) / (i - 4)
                 break
 
             input_dict = dict(input_ids=token_ix)
-    print()
-    return output
-    
-import streamlit as st
+    return output, batch_time
+
+
+tokenizer, model = load_model()
+text = st.text_area("Prefix")
+batch = st.number_input("Variants", value=1)
+
+t = st.empty()
+firstline = ""
+
+def PrintSome(text):
+    global t, firstline
+    firstline += text
+    t.markdown(f"## {firstline}...")
 
+choices, batch_time = PrintContinuation(text, model, PrintSome, batch, 50)
 
-text = st.text_area("Prompt")
-PrintContinuation(text, qmodel, lambda x: t.markdown(f"## {x}..."))
+t.markdown("  \n\n".join(choices) + f"  \n\nBatch:Seconds per batch: {batch_time}, Batch: {batch}")