insert

app.py

@@ -30,6 +30,7 @@ def complete_fn(inpt, max_tokens, min_tokens, alpha_f, alpha_p):
         text = inpt
         counts = [0]*tokenizer.get_vocab_size()
         tokens = tokenizer.encode(inpt).ids
+        yield (None, gr.Text.update(visible=False))
         # yield ("Preproc...", gr.Text.update(visible=False))
         # logits = model.forward(tokens, state)
         for i in range(len(tokens) - 1):
@@ -47,11 +48,11 @@ def complete_fn(inpt, max_tokens, min_tokens, alpha_f, alpha_p):
             counts[token] += 1
             if token == 0:
                 break
-            if i == max_tokens - 1:
-                break
             tokens += [token]
             text = tokenizer.decode(tokens)
             yield (text, None)
+            if i == max_tokens - 1:
+                break
             logits = model.forward_token(token, state)
         yield (text, None)
     except Exception as e:
@@ -60,6 +61,49 @@ def complete_fn(inpt, max_tokens, min_tokens, alpha_f, alpha_p):
     # finally:
     #     return (None, None)
 
+def insert_fn(inpt: str, max_tokens, min_tokens, alpha_f, alpha_p, num_tokens_insert):
+    try:
+        if inpt.count("<|INSERT|>") != 1:
+            yield ("Error...", gr.Text.update(value="Exactly one replace is allowed!", visible=True))
+            return
+        state = rwkv_rs.State(model)
+        text, end = inpt.split("<|INSERT|>")
+        counts = [0]*tokenizer.get_vocab_size()
+        tokens = tokenizer.encode(text).ids
+        tokens_end = tokenizer.encode(end).ids
+        tokens_i = tokens_end[:num_tokens_insert]
+        ins = [0]*len(tokens_i)
+        yield (None, gr.Text.update(visible=False))
+        for i in range(len(tokens) - 1):
+            model.forward_token_preproc(tokens[i], state)
+            yield (tokenizer.decode(tokens[:i + 1]), None)
+        logits = model.forward_token(tokens[-1], state)
+        yield (text, None)
+        max_tokens = int(max_tokens)
+        for i in range(max_tokens):
+            if i < min_tokens:
+                logits[0] = -100
+            for i in range(len(counts)):
+                logits[i] -= (counts[i]* alpha_f) + (float(counts[i] > 0) * alpha_p)
+            token = np.argmax(logits)
+            counts[token] += 1
+            if token == 0:
+                break
+            tokens += [token]
+            ins = ins[1:] + [token]
+            if ins == tokens_i:
+                tokens += tokens_end[num_tokens_insert:]
+                i = max_tokens - 1 # to break earlier...
+            text = tokenizer.decode(tokens)
+            yield (text, None)
+            if i == max_tokens - 1:
+                break
+            logits = model.forward_token(token, state)
+        yield (text, None)
+    except Exception as e:
+        print(e)
+        yield ("Error...", gr.Text.update(value=str(e), visible=True))
+
 def generator_wrap(l, fn):
     def wrap(*args):
         last_i = list([None] * l)
@@ -82,12 +126,14 @@ with gr.Blocks() as app:
             out = gr.TextArea(label="Output")
         complete = gr.Button("Complete", variant="primary")
         c_stop = gr.Button("Stop", variant="stop", visible=False)
-    with gr.Tab("Insert (WIP)"):
-        gr.Markdown("WIP, use `<|INSERT|>` to indicate a place to replace")
+    with gr.Tab("Insert"):
+        gr.Markdown("Use `<|INSERT|>` to indicate a place to replace, if insert fails - end text won't be concatenated")
         with gr.Row():
             inpt_i = gr.TextArea(label="Input")
             out_i = gr.TextArea(label="Output")
-        insert = gr.Button("Insert")
+        num_tokens_insert = gr.Slider(label="Number of tokens to compare for ending (from the beginning of 2nd part)", minimum=1, maximum=2048, value=1024, step=1)
+        insert = gr.Button("Insert", variant="submit")
+        i_stop = gr.Button("Stop", variant="stop", visible=False)
 
     with gr.Column():
         max_tokens = gr.Slider(label="Max Tokens", minimum=1, maximum=4096, step=1, value=767)
@@ -95,10 +141,11 @@ with gr.Blocks() as app:
         alpha_f = gr.Slider(label="Alpha Frequency", minimum=0, maximum=100, step=0.01)
         alpha_p = gr.Slider(label="Alpha Presence", minimum=0, maximum=100, step=0.01)
 
-    G = [complete, c_stop]
-
-    c = complete.click(generator_wrap(2, complete_fn), [inpt, max_tokens, min_tokens, alpha_f, alpha_p], [out, error_box] + G)
+    c = complete.click(generator_wrap(2, complete_fn), [inpt, max_tokens, min_tokens, alpha_f, alpha_p], [out, error_box, complete, c_stop])
     c_stop.click(lambda: (complete.update(visible=True), c_stop.update(visible=False)), inputs=None, outputs=[complete, c_stop], cancels=[c], queue=False)
 
+    i = insert.click(generator_wrap(2, insert_fn), [inpt_i, max_tokens, min_tokens, alpha_f, alpha_p, num_tokens_insert], [out_i, error_box, insert, i_stop])
+    i_stop.click(lambda: (insert.update(visible=True), i_stop.update(visible=False)), inputs=None, outputs=[insert, i_stop], cancels=[i], queue=False)
+
 app.queue(concurrency_count=2)
 app.launch()