Spaces:
Running
Running
File size: 3,122 Bytes
7b9568f |
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 |
import torch
import torch.nn.functional as F
import gradio as gr
from huggingface_hub import hf_hub_download
from namegenerator import Model, ModelConfig
torch.set_grad_enabled(False)
special_tokens = ["<pad>", "<sos>", "<eos>", "<unk>", "0", "1"]
tokens = special_tokens + list("abcdefghijklmnopqrstuvwxyz")
char_to_idx = {char: idx for idx, char in enumerate(tokens)}
idx_to_char = {idx: char for idx, char in enumerate(tokens)}
hf_hub_download(
"karanravindra/namegenerator", "model.pth", subfolder="model", local_dir="."
)
model = Model(
ModelConfig(
vocab_size=len(tokens),
embedding_dim=48,
num_layers=6,
max_length=24, # not padding to nearest 32 because max length of names is 17 - bump this for `theoretically` better performance
q_heads=12,
kv_heads=4,
m=4,
tie_weights=False,
)
)
model.load_state_dict(
torch.load("model/model.pth", map_location="cpu", weights_only=True)
)
model.eval()
def decode(encoded_name: list[int], strip_special_tokens: bool = True) -> str:
if strip_special_tokens:
encoded_name = [
idx
for idx in encoded_name
if idx
not in [char_to_idx["<sos>"], char_to_idx["<eos>"], char_to_idx["<pad>"]]
]
return "".join([idx_to_char[idx] for idx in encoded_name])
def decode_batch(
encoded_names: torch.Tensor, strip_special_tokens: bool = True
) -> list[str]:
return [
decode(encoded_name.tolist(), strip_special_tokens)
for encoded_name in encoded_names
]
def generate_names(n=16, gender=None, temperature=0.6):
model.eval()
if gender is None:
genders = torch.cat(
[
torch.tensor([[char_to_idx["0"]]]).repeat(n // 2, 1),
torch.tensor([[char_to_idx["1"]]]).repeat(n // 2, 1),
],
dim=0,
)
else:
gender = char_to_idx[str(gender)]
genders = torch.full((n, 1), gender)
start_token = torch.tensor([[char_to_idx["<sos>"]]]).repeat(n, 1)
start_token = torch.cat([start_token, genders], dim=1)
generated = start_token
for _ in range(22):
output = model(generated) / temperature
token = torch.multinomial(F.softmax(output[:, -1], dim=1), 1)
generated = torch.cat([generated, token], dim=1)
if token.all() == char_to_idx["<pad>"]:
break
return decode_batch(generated, strip_special_tokens=True)
def generate_name(gender: str, num_names: int, temperature: float):
names = generate_names(num_names, gender, temperature)
names = [name[1:].capitalize() for name in names]
return "\n".join(names)
demo = gr.Interface(
generate_name,
gr.Radio(["Male", "Female"], label="Sex", type="index"),
gr.TextArea(lines=16, label="Generated Names"),
additional_inputs=[
gr.Number(16, label="Number of Names"),
gr.Slider(0.1, 2, 0.6, label="Temperature", step=0.1),
],
title="Name Generator",
description="Generates names based on sex using a GPT-2 model trained on names.",
)
demo.launch()
|