random input example

app.py

@@ -7,6 +7,7 @@ from dearth_config import DearthConfig
 from dearth_model import DearthForCausalLM
 
 import random
+import time
 
 
 
@@ -37,7 +38,7 @@ def generate(input, num_more_tokens):
     if "vocab_size" not in config:
         config['vocab_size'] = tk.vocab_size
     config["attn_window_size"] = 500
-    print(config)
+    # print(config)
     config = DearthConfig(**config)
     model = DearthForCausalLM(config)
 
@@ -45,12 +46,12 @@ def generate(input, num_more_tokens):
 
 
     num_more_tokens = int(num_more_tokens)
-    print(input)
+    # print(input)
     input = input.strip()
     input_ids = tk.encode(input)
     input_ids = [tk.bos_token_id] + input_ids
     input_ids = torch.tensor(input_ids, dtype=torch.long).view(1, -1)
-    print(input_ids)
+    # print(input_ids)
 
     output_ids = input_ids.squeeze(0).tolist()
     for i in range(num_more_tokens):
@@ -58,7 +59,7 @@ def generate(input, num_more_tokens):
         with torch.no_grad():
             output = model(input)[0]
             last_token_logits = output[0, -1, :]
-            last_token_logits_topk = torch.topk(last_token_logits, k=8, dim=-1)
+            last_token_logits_topk = torch.topk(last_token_logits, k=5, dim=-1)
             probs = torch.softmax(last_token_logits_topk.values, dim=-1)
             new_token = torch.multinomial(probs, num_samples=1).item()
             new_token = last_token_logits_topk.indices[new_token].item()
@@ -66,8 +67,8 @@ def generate(input, num_more_tokens):
             break
         output_ids.append(new_token)
 
-    print(output_ids)
-    print(tk.decode(output_ids))
+    # print(output_ids)
+    # print(tk.decode(output_ids))
     output_ids = output_ids[1:]
 
     return tk.decode(output_ids)
@@ -77,7 +78,7 @@ example_input = ["Once upon a time, there was a little girl",
                  "It was a warm summer day when Billy and",
 ]
 
-
+ui_title = "Tinystories LM 11M"
 Description = """
 This is a small language model with 11M parameters, trained with the TinyStories dataset, and distilled from a 28M parameter teacher model.\n
 This model has been trained with 512M tokens, which is about 0.9 epoch of the TinyStories dataset.\n
@@ -85,17 +86,38 @@ The PPL on the validation set is 1.7, in comparison, the teacher model has a PPL
 """
 
 
-demo = gr.Interface(
-    fn=generate,
-    title="Tinystories LM 11M",
-    description=Description,
-    inputs=[
-        gr.Textbox(lines=5, label="Input Text", value=example_input[random.randint(0, len(example_input)-1)]),
-        gr.Slider(16, 64, step=1.0, value=32, label="more tokens", info="")
-    ],
-    outputs="text"
-)
+# demo = gr.Interface(
+#     fn=generate,
+#     title="Tinystories LM 11M",
+#     description=Description,
+#     inputs=[
+#         gr.Textbox(lines=5, label="Input Text", value=example_input[random.randint(0, len(example_input)-1)]),
+#         gr.Slider(16, 64, step=1.0, value=32, label="more tokens", info="")
+#     ],
+#     outputs="text"
+# )
+
+with open("./random_input_example.js" , "r") as f:
+    file_content = f.read()
 
 if __name__ == "__main__":
+    with gr.Blocks(
+        title="Tinystories LM 11M",
+        js="./random_input_example.js"
+    ) as demo:
+        with gr.Blocks(title="Description"):
+            gr.HTML(f"<h1>{ui_title}</h1>")
+            gr.Markdown(Description)
+        with gr.Row():
+            with gr.Column():
+                inp = gr.Textbox(lines=5, label="Input Text", value=example_input[random.randint(0, len(example_input)-1)], elem_id="input_textbox")
+                generate_max_slider = gr.Slider(16, 64, step=1.0, value=32, label="more tokens", info="")
+                generate_button = gr.Button(value="Generate")
+            with gr.Column():
+                out = gr.Textbox(lines=5, label="Output Text", value="")
+                out.readonly = True
+            @generate_button.click(inputs=[inp, generate_max_slider], outputs=[out])
+            def generate_inside(input, num_more_tokens):
+                return generate(input, num_more_tokens)
     demo.queue()
-    demo.launch(show_api=False)
+    demo.launch()

random_input_example.js

@@ -0,0 +1,8 @@
+() => {
+    let example_input = ["Once upon a time, there was a little girl", 
+                 "John and Sarah were playing together in their backyard when",
+                 "It was a warm summer day when Billy and",
+    ];
+    let input_textbox = document.querySelector("#input_textbox textarea")
+    input_textbox.value = example_input[Math.floor(Math.random() * example_input.length)];
+}