Upload folder using huggingface_hub

inference.py

@@ -0,0 +1,108 @@
+import torch
+import time
+from tokenizers import Tokenizer
+from miniGPT import MiniGPT
+
+# --- 1. Load tokenizer and model ---
+tokenizer = Tokenizer.from_file("wordlevel.json")
+vocab_size = tokenizer.get_vocab_size()
+
+# Set model parameters to match your trained model
+model = MiniGPT(
+    vocab_size=vocab_size,
+    embed_dim=128,
+    num_heads=4,
+    ff_dim=512,
+    num_layers=4,
+    max_seq_len=128
+)
+checkpoint_path = "model_checkpoint_step20000.pt"
+checkpoint = torch.load(checkpoint_path, map_location='cpu')
+model.load_state_dict(checkpoint['model_state_dict'])
+model.eval()
+
+# --- 2. Show model parameter count ---
+num_params = sum(p.numel() for p in model.parameters())
+print(f"Model parameters: {num_params:,}")
+
+# --- 3. Sampling helpers ---
+
+def top_k_logits(logits, k):
+    """Keep only top-k tokens with highest probability."""
+    values, _ = torch.topk(logits, k)
+    min_values = values[:, -1].unsqueeze(1)
+    logits[logits < min_values] = -float('Inf')
+    return logits
+
+def top_p_logits(logits, p=0.9):
+    """Keep the smallest set of tokens with cumulative probability >= p."""
+    sorted_logits, sorted_indices = torch.sort(logits, descending=True, dim=-1)
+    cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+    
+    sorted_indices_to_remove = cumulative_probs > p
+    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+    sorted_indices_to_remove[..., 0] = 0
+
+    for batch in range(logits.size(0)):
+        remove_ids = sorted_indices[batch][sorted_indices_to_remove[batch]]
+        logits[batch, remove_ids] = -float('Inf')
+    
+    return logits
+
+# --- 4. Streaming generation function ---
+def generate_stream(
+    model, tokenizer, prompt, 
+    max_new_tokens=50, 
+    temperature=1.0, 
+    top_k=None, 
+    top_p=None,
+    repetition_penalty=2.0
+):
+    idx = torch.tensor([tokenizer.encode(prompt).ids], dtype=torch.long)
+    generated = []
+    start_time = time.time()
+
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            if idx.shape[1] >= model.max_seq_len:
+                break
+
+            logits = model(idx)
+            logits = logits[:, -1, :] / temperature
+
+            # Apply repetition penalty
+            for token_id in set(generated):
+                logits[0, token_id] /= repetition_penalty
+
+            # Apply Top-K and/or Top-P filtering
+            if top_k is not None:
+                logits = top_k_logits(logits, top_k)
+            if top_p is not None:
+                logits = top_p_logits(logits, top_p)
+
+            probs = torch.softmax(logits, dim=-1)
+            next_id = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat([idx, next_id], dim=1)
+            generated.append(next_id.item())
+            print(tokenizer.decode([next_id.item()]), end=' ', flush=True)
+
+    elapsed = time.time() - start_time
+    tps = len(generated) / elapsed if elapsed > 0 else 0
+    print(f"\n[Generated {len(generated)} tokens in {elapsed:.2f} seconds | {tps:.2f} tokens/sec]")
+    return idx
+
+# --- 5. Main input loop ---
+while True:
+    prompt = input("\nEnter your prompt (or type 'exit' to quit): ")
+    if prompt.lower() == 'exit':
+        break
+
+    print("\nStreaming output:")
+    generate_stream(
+        model, tokenizer, prompt, 
+        max_new_tokens=90,
+        temperature=2.0,
+        top_k=100,
+        top_p=0.9,
+        repetition_penalty=1.8
+    )

miniGPT.py

@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformer import TransformerBlock
+
+class MiniGPT(nn.Module):
+    def __init__(self, vocab_size, embed_dim, num_heads, ff_dim, num_layers, max_seq_len):
+        super().__init__()
+        self.max_seq_len = max_seq_len 
+        self.token_embedding = nn.Embedding(vocab_size, embed_dim)
+        self.pos_embedding = nn.Embedding(max_seq_len, embed_dim)
+        self.blocks = nn.Sequential(
+            *[TransformerBlock(embed_dim, num_heads, ff_dim) for _ in range(num_layers)]
+        )
+        self.ln_f = nn.LayerNorm(embed_dim)
+        self.head = nn.Linear(embed_dim, vocab_size, bias=False)
+
+        self.head.weight = self.token_embedding.weight
+
+    def forward(self, idx, mask=None):
+        B, T = idx.shape
+        tok_emb = self.token_embedding(idx)
+        pos = torch.arange(T,device=idx.device).unsqueeze(0)
+        pos_emb = self.pos_embedding(pos)
+        x = tok_emb + pos_emb
+        x = self.blocks(x, mask=mask) if mask is not None else self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.head(x)
+        return logits
+    

multiheadattention.py

@@ -0,0 +1,34 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, embed_dim, num_heads):
+        super().__init__()
+        assert embed_dim % num_heads == 0, "Embedding dim must be divisible by num heads"
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.head_dim = embed_dim // num_heads
+
+        self.qkv_proj = nn.Linear(embed_dim, 3 * embed_dim)
+        self.out_proj = nn.Linear(embed_dim, embed_dim)
+
+    def forward(self, x, mask=None):
+        B, T, C = x.shape 
+        qkv = self.qkv_proj(x)
+        qkv = qkv.reshape(B, T, self.num_heads, 3 * self.head_dim)
+        qkv = qkv.permute(0, 2, 1, 3)
+        q, k, v = qkv.chunk(3, dim=-1)
+
+        attn_scores = (q @ k.transpose(-2, -1)) / (self.head_dim ** 0.5)
+
+        if mask is not None:
+            attn_scores = attn_scores.masked_fill(mask == 0, float('-inf'))
+
+        attn_weights = F.softmax(attn_scores, dim=-1)
+        attn_output = attn_weights @ v
+
+
+        attn_output = attn_output.transpose(1, 2).reshape(B, T, C)
+        ouptut = self.out_proj(attn_output)
+        return ouptut

transformer.py

@@ -0,0 +1,24 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from multiheadattention import MultiHeadAttention
+
+class TransformerBlock(nn.Module):
+    def __init__(self, embed_dim, num_heads, ff_dim):
+        super().__init__()
+        self.attn = MultiHeadAttention(embed_dim, num_heads)
+        self.ln1 = nn.LayerNorm(embed_dim)
+        self.ff = nn.Sequential(
+            nn.Linear(embed_dim, ff_dim),
+            nn.GELU(),
+            nn.Linear(ff_dim, embed_dim)
+        )
+
+        self.ln2 = nn.LayerNorm(embed_dim)
+
+    def forward(self, x, mask=None):
+        x = x + self.attn(self.ln1(x), mask = mask)
+        x = x + self.ff(self.ln2(x))
+        return x
+
+