make space

.gitignore

@@ -0,0 +1,2 @@
+.venv
+model

app.py

@@ -0,0 +1,104 @@
+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()

requirements.txt

@@ -0,0 +1,68 @@
+namegenerator @ git+https://github.com/karanravindra/namegenerator@main
+gradio >= 5.0.0
+huggingface_hub >= 0.26.0
+
+# This file was autogenerated via `uv export`.
+appnope==0.1.4 ; platform_system == 'Darwin'
+asttokens==2.4.1
+cffi==1.17.1 ; implementation_name == 'pypy'
+colorama==0.4.6 ; sys_platform == 'win32' or platform_system == 'Windows'
+comm==0.2.2
+debugpy==1.8.7
+decorator==5.1.1
+einops==0.8.0
+exceptiongroup==1.2.2 ; python_full_version < '3.11'
+executing==2.1.0
+filelock==3.16.1
+fsspec==2024.10.0
+ipykernel==6.29.5
+ipython==8.29.0
+ipywidgets==8.1.5
+jedi==0.19.1
+jupyter-client==8.6.3
+jupyter-core==5.7.2
+jupyterlab-widgets==3.0.13
+matplotlib-inline==0.1.7
+mpmath==1.3.0
+nest-asyncio==1.6.0
+networkx==3.4.2
+numpy==2.1.2
+nvidia-cublas-cu12==12.4.5.8 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-cupti-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-nvrtc-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cuda-runtime-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cudnn-cu12==9.1.0.70 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cufft-cu12==11.2.1.3 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-curand-cu12==10.3.5.147 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cusolver-cu12==11.6.1.9 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-cusparse-cu12==12.3.1.170 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nccl-cu12==2.21.5 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nvjitlink-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+nvidia-nvtx-cu12==12.4.127 ; platform_machine == 'x86_64' and platform_system == 'Linux'
+packaging==24.1
+parso==0.8.4
+pexpect==4.9.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'
+platformdirs==4.3.6
+polars==1.12.0
+prompt-toolkit==3.0.48
+psutil==6.1.0
+ptyprocess==0.7.0 ; sys_platform != 'emscripten' and sys_platform != 'win32'
+pure-eval==0.2.3
+pycparser==2.22 ; implementation_name == 'pypy'
+pygments==2.18.0
+python-dateutil==2.9.0.post0
+pywin32==308 ; platform_python_implementation != 'PyPy' and sys_platform == 'win32'
+pyzmq==26.2.0
+setuptools==75.3.0 ; python_full_version >= '3.12'
+six==1.16.0
+stack-data==0.6.3
+sympy==1.13.1
+torch==2.5.0
+torchinfo==1.8.0
+tornado==6.4.1
+tqdm==4.66.6
+traitlets==5.14.3
+triton==3.1.0 ; python_full_version < '3.13' and platform_machine == 'x86_64' and platform_system == 'Linux'
+typing-extensions==4.12.2
+wcwidth==0.2.13
+widgetsnbextension==4.0.13