Upload exported SimpleStories Hindi 10M/11M model with safetensors, custom modeling.py and tokenizer
70c8597 verified | """ | |
| SimpleStories Hindi - Custom HuggingFace Model | |
| A small LLaMA-style decoder-only Transformer trained from scratch on Hindi stories. | |
| Architecture: RoPE + SwiGLU + RMSNorm + FlashAttention (via scaled_dot_product_attention) | |
| Usage: | |
| from transformers import AutoModelForCausalLM, AutoTokenizer | |
| tokenizer = AutoTokenizer.from_pretrained("your-repo", trust_remote_code=True) | |
| model = AutoModelForCausalLM.from_pretrained("your-repo", trust_remote_code=True) | |
| inputs = tokenizer("एक समय की बात है", return_tensors="pt") | |
| outputs = model.generate(**inputs, max_new_tokens=100, do_sample=True, temperature=0.8) | |
| print(tokenizer.decode(outputs[0], skip_special_tokens=True)) | |
| """ | |
| import math | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from transformers import PreTrainedModel, PretrainedConfig, GenerationMixin | |
| from transformers.modeling_outputs import CausalLMOutputWithPast | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Config | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class SimpleStoriesConfig(PretrainedConfig): | |
| """Configuration for the SimpleStories Hindi model.""" | |
| model_type = "simple_stories" | |
| def __init__( | |
| self, | |
| vocab_size=6000, | |
| max_seq_len=512, | |
| d_model=384, | |
| n_layers=10, | |
| n_heads=8, | |
| multiple_of=64, | |
| norm_eps=1e-5, | |
| dropout=0.0, | |
| pad_token_id=0, | |
| unk_token_id=1, | |
| bos_token_id=2, | |
| eos_token_id=3, | |
| **kwargs | |
| ): | |
| super().__init__( | |
| pad_token_id=pad_token_id, | |
| unk_token_id=unk_token_id, | |
| bos_token_id=bos_token_id, | |
| eos_token_id=eos_token_id, | |
| **kwargs | |
| ) | |
| self.vocab_size = vocab_size | |
| self.max_seq_len = max_seq_len | |
| self.d_model = d_model | |
| self.n_layers = n_layers | |
| self.n_heads = n_heads | |
| self.multiple_of = multiple_of | |
| self.norm_eps = norm_eps | |
| self.dropout = dropout | |
| # HuggingFace standard aliases (needed by DynamicCache, generation, etc.) | |
| self.num_hidden_layers = n_layers | |
| self.hidden_size = d_model | |
| self.num_attention_heads = n_heads | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # RoPE helpers | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| def _rotate_half(x: torch.Tensor) -> torch.Tensor: | |
| x1 = x[..., : x.shape[-1] // 2] | |
| x2 = x[..., x.shape[-1] // 2 :] | |
| return torch.cat((-x2, x1), dim=-1) | |
| def _precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0): | |
| freqs = 1.0 / (theta ** (torch.arange(0, dim, 2).float() / dim)) | |
| t = torch.arange(end, dtype=torch.float32) | |
| freqs = torch.outer(t, freqs) | |
| cos = torch.cos(freqs) | |
| sin = torch.sin(freqs) | |
| # duplicate so shape is [end, dim] instead of [end, dim//2] | |
| cos = torch.cat([cos, cos], dim=-1) | |
| sin = torch.cat([sin, sin], dim=-1) | |
| return cos, sin | |
| def _apply_rotary_emb(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| T = x.shape[2] | |
| cos_t = cos[:T, :].unsqueeze(0).unsqueeze(1) # [1, 1, T, head_dim] | |
| sin_t = sin[:T, :].unsqueeze(0).unsqueeze(1) | |
| return (x * cos_t) + (_rotate_half(x) * sin_t) | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Sub-modules | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class RMSNorm(nn.Module): | |
| def __init__(self, dim: int, eps: float = 1e-5): | |
| super().__init__() | |
| self.eps = eps | |
| self.weight = nn.Parameter(torch.ones(dim)) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| variance = x.pow(2).mean(-1, keepdim=True) | |
| return x * torch.rsqrt(variance + self.eps) * self.weight | |
| class FeedForward(nn.Module): | |
| """SwiGLU feed-forward block.""" | |
| def __init__(self, dim: int, hidden_dim: int = None, multiple_of: int = 32, dropout: float = 0.0): | |
| super().__init__() | |
| if hidden_dim is None: | |
| hidden_dim = int(2 * 4 * dim / 3) | |
| hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of) | |
| self.w1 = nn.Linear(dim, hidden_dim, bias=False) | |
| self.w2 = nn.Linear(dim, hidden_dim, bias=False) | |
| self.w3 = nn.Linear(hidden_dim, dim, bias=False) | |
| self.dropout = nn.Dropout(dropout) if dropout > 0.0 else None | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| out = F.silu(self.w1(x)) * self.w2(x) | |
| out = self.w3(out) | |
| if self.dropout is not None: | |
| out = self.dropout(out) | |
| return out | |
| class Attention(nn.Module): | |
| def __init__(self, config: SimpleStoriesConfig): | |
| super().__init__() | |
| self.n_heads = config.n_heads | |
| self.d_model = config.d_model | |
| self.head_dim = config.d_model // config.n_heads | |
| self.wq = nn.Linear(config.d_model, config.d_model, bias=False) | |
| self.wk = nn.Linear(config.d_model, config.d_model, bias=False) | |
| self.wv = nn.Linear(config.d_model, config.d_model, bias=False) | |
| self.wo = nn.Linear(config.d_model, config.d_model, bias=False) | |
| self.dropout_p = config.dropout | |
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| B, T, C = x.shape | |
| q = self.wq(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) | |
| k = self.wk(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) | |
| v = self.wv(x).view(B, T, self.n_heads, self.head_dim).transpose(1, 2) | |
| q = _apply_rotary_emb(q, cos, sin) | |
| k = _apply_rotary_emb(k, cos, sin) | |
| dropout_p = self.dropout_p if self.training else 0.0 | |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=dropout_p, is_causal=True) | |
| out = out.transpose(1, 2).contiguous().view(B, T, C) | |
| return self.wo(out) | |
| class TransformerBlock(nn.Module): | |
| def __init__(self, config: SimpleStoriesConfig): | |
| super().__init__() | |
| self.attention = Attention(config) | |
| self.feed_forward = FeedForward(dim=config.d_model, multiple_of=config.multiple_of, dropout=config.dropout) | |
| self.attention_norm = RMSNorm(config.d_model, eps=config.norm_eps) | |
| self.ffn_norm = RMSNorm(config.d_model, eps=config.norm_eps) | |
| def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor: | |
| x = x + self.attention(self.attention_norm(x), cos, sin) | |
| x = x + self.feed_forward(self.ffn_norm(x)) | |
| return x | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| # Main model | |
| # ────────────────────────────────────────────────────────────────────────────── | |
| class SimpleStoriesForCausalLM(PreTrainedModel, GenerationMixin): | |
| """LLaMA-style causal language model for Hindi story generation.""" | |
| config_class = SimpleStoriesConfig | |
| _tied_weights_keys = ["output.weight"] | |
| def __init__(self, config: SimpleStoriesConfig): | |
| super().__init__(config) | |
| self.config = config | |
| self.tok_embeddings = nn.Embedding(config.vocab_size, config.d_model) | |
| self.dropout = nn.Dropout(config.dropout) if config.dropout > 0.0 else None | |
| self.layers = nn.ModuleList([TransformerBlock(config) for _ in range(config.n_layers)]) | |
| self.norm = RMSNorm(config.d_model, eps=config.norm_eps) | |
| self.output = nn.Linear(config.d_model, config.vocab_size, bias=False) | |
| # weight tying | |
| self.tok_embeddings.weight = self.output.weight | |
| self.post_init() | |
| # Required by PreTrainedModel — we skip re-init since we load trained weights. | |
| def _init_weights(self, module): | |
| pass | |
| def forward( | |
| self, | |
| input_ids: torch.LongTensor = None, | |
| attention_mask: torch.Tensor = None, | |
| labels: torch.LongTensor = None, | |
| past_key_values=None, | |
| use_cache: bool = False, | |
| return_dict: bool = True, | |
| **kwargs, | |
| ): | |
| B, T = input_ids.shape | |
| x = self.tok_embeddings(input_ids) | |
| if self.dropout is not None: | |
| x = self.dropout(x) | |
| # Compute RoPE embeddings on the fly (avoids meta-device buffer issues) | |
| head_dim = self.config.d_model // self.config.n_heads | |
| cos, sin = _precompute_freqs_cis(dim=head_dim, end=T) | |
| cos = cos.to(x.device, dtype=x.dtype) | |
| sin = sin.to(x.device, dtype=x.dtype) | |
| for layer in self.layers: | |
| x = layer(x, cos, sin) | |
| x = self.norm(x) | |
| logits = self.output(x) | |
| loss = None | |
| if labels is not None: | |
| loss = F.cross_entropy( | |
| logits.view(-1, logits.size(-1)), | |
| labels.view(-1), | |
| ignore_index=-100, | |
| ) | |
| if not return_dict: | |
| return (loss, logits) if loss is not None else (logits,) | |
| return CausalLMOutputWithPast(loss=loss, logits=logits) | |
| def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs): | |
| return {"input_ids": input_ids, "use_cache": False} | |