Spaces:
Sleeping
Sleeping
File size: 4,661 Bytes
4ae913a 7569c73 4ae913a ce226d1 4ae913a ce226d1 7569c73 ce226d1 4ae913a 7569c73 4ae913a 7569c73 4ae913a 7569c73 4ae913a 7569c73 4ae913a 7569c73 4ae913a 7569c73 4ae913a ce226d1 7569c73 ce226d1 7569c73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 |
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() |