blergh

app.py

@@ -104,6 +104,64 @@ def insert_fn(inpt: str, max_tokens, min_tokens, alpha_f, alpha_p, num_tokens_in
         print(e)
         yield ("Error...", gr.Text.update(value=str(e), visible=True))
 
+# def classify_fn_inner(inpt, clas):
+#     state = rwkv_rs.State(model)
+#     tokens = tokenizer.encode(f"This is an example of {clas} text: {inpt}").ids
+#     for i in tokens[:-2]:
+#         model.forward_token_preproc(i, state)
+#     # state_2 = state.copy()
+
+#     logit_x_1 = softmax(model.forward_token(tokens[-2], state))
+#     logit_y_1 = softmax(model.forward_token(tokens[-1], state))
+#     # shapep = logit_x_1.shape[0] * 0.9
+#     # s = np.sort(logit_y_1)[::-1]
+#     # c = s[np.argmax(np.cumsum(s) > 0.9)]
+#     # logit_y_1[logit_y_1 < c] = 0
+#     loss_1 = -np.sum(logit_y_1 * np.log(logit_x_1)) / logit_x_1.shape[0]
+
+#     # I forgor that I do not return the preproc shit...
+#     # logit_x_2 = model.forward_token_preproc(tokens[-2], state_2)
+#     # logit_y_2 = model.forward_token_preproc(tokens[-1], state_2)
+#     return (loss_1, None)
+
+def classify_fn_inner2(inpt, clas):
+    state = rwkv_rs.State(model)
+    tokens = tokenizer.encode(f"This is an example of {clas} text:").ids
+    for i in tokens:
+        model.forward_token_preproc(i, state)
+
+    logits = []
+    tokens = tokenizer.encode(f" {inpt}\n").ids
+    for i in tokens[:-1]:
+        logits.append(model.forward_token(i, state))
+    logit_x = [softmax(i) for i in logits]
+    loss = -np.sum([ x[y] for x, y in zip(logit_x, tokens[1:]) ]) / len(logit_x)
+
+    return loss 
+
+def softmax(x):
+    e = np.exp(x - np.max(x))
+    return e / e.sum()
+
+# TODO: maybe make a function with pos/neg inputs?
+def classify_fn(inpt, clas, clasneg):
+    # loss_1, loss_2 = classify_fn_inner(inpt, clas)
+    # loss_1_neg, loss_2_neg = classify_fn_inner(inpt, clasneg)
+
+    # print(loss_1, loss_1_neg, end=' | ')
+    # # We negate the loss because we want to know who's closer to 0
+    # loss_1, loss_1_neg = softmax([-loss_1, -loss_1_neg])
+    # print(loss_1, loss_1_neg)
+
+    loss_3 = classify_fn_inner2(inpt, clas)
+    loss_3_neg = classify_fn_inner2(inpt, clasneg)
+    print(loss_3, loss_3_neg, end=' | ')
+    loss_3, loss_3_neg = softmax([-loss_3, -loss_3_neg])
+    print(loss_3, loss_3_neg)
+
+    # return ({"v1_pos": loss_1, "v1_neg": loss_1_neg, "v3_pos": loss_3, "v3_neg": loss_3_neg})
+    return ({"+": loss_3, "-": loss_3_neg})
+
 def generator_wrap(l, fn):
     def wrap(*args):
         last_i = list([None] * l)
@@ -134,6 +192,14 @@ with gr.Blocks() as app:
         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="primary")
         i_stop = gr.Button("Stop", variant="stop", visible=False)
+    with gr.Tab("Classification W/O head"):
+        gr.Markdown("This is an experimental classification attempt based on [this Twitter post](https://twitter.com/aicrumb/status/1625239547268280321)\n\nSettings at the bottom do no affect this example.")
+        with gr.Row():
+            inpt_c = gr.TextArea(label="Input")
+            out_c = gr.Label(label="Output")
+        clas = gr.Textbox(label="+ NL class/example to check against.")
+        clasneg = gr.Textbox(label="- NL class/example to check against.")
+        classify = gr.Button("Classify", variant="primary")
 
     with gr.Column():
         max_tokens = gr.Slider(label="Max Tokens", minimum=1, maximum=4096, step=1, value=767)
@@ -147,5 +213,7 @@ with gr.Blocks() as app:
     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)
 
+    classify.click(classify_fn, [inpt_c, clas, clasneg], [out_c])
+
 app.queue(concurrency_count=2)
 app.launch()