app.py

import gradio as gr
import transformers
import torch
import yaml

from dearth_config import DearthConfig
from dearth_model import DearthForCausalLM

import random
import time



tk = transformers.AutoTokenizer.from_pretrained("./tk")
model_path = "./ts100-re2-h1-4000-model.pt"
states = torch.load(model_path, map_location="cpu")
model_states = states
unwanted_prefix_dueto_compile = '_orig_mod.'
unwanted_prefix_dueto_ddp = 'module.'
unwanted_prefix_dueto_ddp_compiled = 'module._orig_mod.'

for k,v in list(model_states.items()):
    if k.startswith(unwanted_prefix_dueto_ddp_compiled):
        new_key = k[len(unwanted_prefix_dueto_ddp_compiled):]
        model_states[k[len(unwanted_prefix_dueto_ddp_compiled):]] = model_states.pop(k)
    elif k.startswith(unwanted_prefix_dueto_ddp):
        new_key = k[len(unwanted_prefix_dueto_ddp):]
        model_states[k[len(unwanted_prefix_dueto_ddp):]] = model_states.pop(k)
    elif k.startswith(unwanted_prefix_dueto_compile):
        new_key = k[len(unwanted_prefix_dueto_compile):]
        model_states[k[len(unwanted_prefix_dueto_compile):]] = model_states.pop(k)

def generate(input, num_more_tokens):
  
    yml_path = "./ts100-re2-h1.yml"
    with open(yml_path, "r") as f:
        config = yaml.load(f, Loader=yaml.FullLoader)['model']
    if "vocab_size" not in config:
        config['vocab_size'] = tk.vocab_size
    config["attn_window_size"] = 500
    # print(config)
    config = DearthConfig(**config)
    model = DearthForCausalLM(config)

    model.load_state_dict(model_states)


    num_more_tokens = int(num_more_tokens)
    # 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)

    output_ids = input_ids.squeeze(0).tolist()
    for i in range(num_more_tokens):
        input = torch.tensor(output_ids, dtype=torch.long).view(1, -1)
        with torch.no_grad():
            output = model(input)[0]
            last_token_logits = output[0, -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()
        if new_token == tk.eos_token_id:
            break
        output_ids.append(new_token)

    # print(output_ids)
    # print(tk.decode(output_ids))
    output_ids = output_ids[1:]

    return tk.decode(output_ids)

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",
]

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
The PPL on the validation set is 1.7, in comparison, the teacher model has a PPL of 0.9. Lower PPL means better performance.\n
"""


# 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()