Spaces:
Sleeping
Sleeping
| import os | |
| import gradio as gr | |
| import torch | |
| import torch.nn.functional as F | |
| try: | |
| import spaces # available on Hugging Face ZeroGPU Spaces | |
| gpu = spaces.GPU | |
| except ImportError: # running locally | |
| gpu = lambda fn: fn | |
| class Sequential: | |
| def __init__(self, layers): | |
| self.layers = layers | |
| def __call__(self, x): | |
| for layer in self.layers: | |
| x = layer(x) | |
| self.out = x | |
| return self.out | |
| class Embedding: | |
| def __init__(self, num_embeddings, embeddings_size): | |
| self.weight = torch.randn((num_embeddings, embeddings_size)) | |
| def __call__(self, ix): | |
| self.out = self.weight[ix] | |
| return self.out | |
| class Flatten: | |
| def __init__(self, n): | |
| self.n = n | |
| def __call__(self, x): | |
| batch_size, block_size, n_emb = x.shape | |
| x = x.view(batch_size, block_size // self.n, n_emb * self.n) | |
| if x.shape[1] == 1: | |
| x = x.squeeze(1) | |
| self.out = x | |
| return self.out | |
| class Linear: | |
| def __init__(self, fanin, fanout, bias=True): | |
| self.weight = torch.randn((fanin, fanout)) / (fanin**0.5) | |
| self.bias = torch.randn(fanout) if bias else None | |
| def __call__(self, x): | |
| self.out = x @ self.weight | |
| if self.bias is not None: | |
| self.out += self.bias | |
| return self.out | |
| class BatchNorm1d: | |
| def __init__(self, dim, eps=1e-5, momentum=0.1): | |
| self.eps = eps | |
| self.momentum = momentum | |
| self.training = True | |
| self.gamma = torch.ones(dim) | |
| self.beta = torch.zeros(dim) | |
| self.running_mean = torch.zeros(dim) | |
| self.running_var = torch.ones(dim) | |
| def __call__(self, x): | |
| # inference only: always use running stats | |
| xhat = (x - self.running_mean) / torch.sqrt(self.running_var + self.eps) | |
| self.out = self.gamma * xhat + self.beta | |
| return self.out | |
| class Tanh: | |
| def __call__(self, x): | |
| self.out = torch.tanh(x) | |
| return self.out | |
| model_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "wavenet.pt") | |
| print("model path:", model_path) | |
| ckpt = torch.load(model_path, map_location="cpu") | |
| itos = ckpt["itos"] | |
| stoi = {s: i for i, s in itos.items()} | |
| block_size = ckpt["block_size"] | |
| vocab_size = len(itos) | |
| # Use the SAME hyperparameters and layer order as in training | |
| n_emb = 50 | |
| n_hidden = 300 | |
| model = Sequential([ | |
| Embedding(vocab_size, n_emb), | |
| Flatten(2), Linear(n_emb * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), | |
| Flatten(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), | |
| Flatten(2), Linear(n_hidden * 2, n_hidden, bias=False), BatchNorm1d(n_hidden), Tanh(), | |
| Linear(n_hidden, vocab_size), | |
| ]) | |
| for key, t in ckpt["state"].items(): | |
| i, name = key.split(".") | |
| setattr(model.layers[int(i)], name, t) | |
| def generate_one(prefix="", temperature=1.0, max_len=30): | |
| context = [0] * block_size | |
| out = [] | |
| for ch in prefix: | |
| context = context[1:] + [stoi[ch]] | |
| out.append(stoi[ch]) | |
| while len(out) < max_len: | |
| logits = model(torch.tensor([context])) | |
| probs = F.softmax(logits / temperature, dim=1) | |
| ix = torch.multinomial(probs, num_samples=1).item() | |
| if ix == 0: # end-of-name token | |
| break | |
| context = context[1:] + [ix] | |
| out.append(ix) | |
| return "".join(itos[i] for i in out) | |
| # ZeroGPU Spaces require at least one @spaces.GPU function; the model itself runs on CPU | |
| def generate_names(num_names, prefix, temperature): | |
| prefix = (prefix or "").strip().lower() | |
| unknown = sorted({ch for ch in prefix if ch not in stoi or stoi[ch] == 0}) | |
| if unknown: | |
| raise gr.Error(f"Unsupported characters in prefix: {' '.join(unknown)}") | |
| names = [generate_one(prefix, temperature).capitalize() for _ in range(int(num_names))] | |
| return "\n".join(names) | |
| with gr.Blocks() as iface: | |
| gr.Markdown( | |
| """ | |
| # 🇮🇳 Indian Name Generator | |
| A character-level language model trained on Indian names, based on the **WaveNet architecture**. | |
| It generates new, name-like words one character at a time. | |
| ### Model details: | |
| - **All layers are custom built** from scratch in PyTorch (Embedding, Flatten, Linear, | |
| BatchNorm1d, Tanh, Sequential) — no `torch.nn` layers used | |
| - **WaveNet architecture**: characters are fused hierarchically, 2 at a time, over a | |
| context of 8 previous characters | |
| - **Embedding size (n_emb):** 50 | |
| - **Hidden units (n_hidden):** 300 | |
| - **Total parameters:** 401,277 | |
| - **Dataset:** [swami93/indian-names-1.5M](https://huggingface.co/datasets/swami93/indian-names-1.5M) | |
| - **Training notebook:** [Kaggle — predict-indian-names-wavenet-arch](https://www.kaggle.com/code/eluriadarsh/predict-indian-names-wavenet-arch) | |
| ### Performance (cross-entropy loss): | |
| | Split | Loss | | |
| |---|---| | |
| | Training | 1.90 | | |
| | Validation | 2.01 | | |
| | Test | 2.01 | | |
| ### How to use: | |
| - Choose how many names to generate | |
| - Optionally give a starting prefix (e.g. `pra`, `su`, `ka`) | |
| - **Temperature**: lower = safer, more common-sounding names; higher = more creative | |
| """ | |
| ) | |
| with gr.Row(): | |
| with gr.Column(): | |
| num_names = gr.Slider(1, 50, value=10, step=1, label="Number of names") | |
| prefix = gr.Textbox(label="Starts with (optional)", placeholder="e.g. pra") | |
| temperature = gr.Slider(0.3, 2.0, value=1.0, step=0.1, label="Temperature") | |
| submit_btn = gr.Button("✨ Generate Names", variant="primary") | |
| with gr.Column(): | |
| output = gr.Textbox(label="Generated names", lines=15) | |
| gr.Examples( | |
| examples=[[10, "", 1.0], [10, "pra", 0.8], [10, "su", 1.0], [10, "ka", 1.3]], | |
| inputs=[num_names, prefix, temperature], | |
| label="Try an example", | |
| ) | |
| submit_btn.click(fn=generate_names, inputs=[num_names, prefix, temperature], outputs=output) | |
| if __name__ == "__main__": | |
| iface.launch(server_name="0.0.0.0", server_port=7860) | |