Upload chat.py with huggingface_hub

chat.py

@@ -0,0 +1,170 @@
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+import sys
+import re
+from indic_transliteration import sanscript
+from indic_transliteration.sanscript import transliterate
+
+# 1. COPY HYPERPARAMETERS & MODEL ARCHITECTURE
+# (Must match train_llm.py EXACTLY)
+block_size = 256
+n_embd = 384
+n_head = 6
+n_layer = 6
+dropout = 0.2
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+
+# Load the dataset to rebuild the vocabulary mapping
+with open('dataset_awadhi_tulsidas_verified.txt', 'r', encoding='utf-8') as f:
+    text = f.read()
+
+chars = sorted(list(set(text)))
+vocab_size = len(chars)
+
+stoi = { ch:i for i,ch in enumerate(chars) }
+itos = { i:ch for i,ch in enumerate(chars) }
+def encode(s):
+    unknown = [c for c in s if c not in stoi]
+    if unknown:
+        raise ValueError(f"Unsupported character(s) in Devanagari output: {unknown}. Make sure you only use standard Hindi characters.")
+    return [stoi[c] for c in s]
+decode = lambda l: ''.join([itos[i] for i in l])
+
+class Head(nn.Module):
+    def __init__(self, head_size):
+        super().__init__()
+        self.key   = nn.Linear(n_embd, head_size, bias=False)
+        self.query = nn.Linear(n_embd, head_size, bias=False)
+        self.value = nn.Linear(n_embd, head_size, bias=False)
+        self.register_buffer('tril', torch.tril(torch.ones(block_size, block_size)))
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        B, T, C = x.shape
+        k = self.key(x)   
+        q = self.query(x) 
+        wei = q @ k.transpose(-2, -1) * (k.shape[-1] ** -0.5)
+        wei = wei.masked_fill(self.tril[:T, :T] == 0, float('-inf'))
+        wei = F.softmax(wei, dim=-1)
+        wei = self.dropout(wei)
+        v   = self.value(x)
+        out = wei @ v      
+        return out
+
+class MultiHeadAttention(nn.Module):
+    def __init__(self, num_heads, head_size):
+        super().__init__()
+        self.heads   = nn.ModuleList([Head(head_size) for _ in range(num_heads)])
+        self.proj    = nn.Linear(n_embd, n_embd)
+        self.dropout = nn.Dropout(dropout)
+
+    def forward(self, x):
+        out = torch.cat([h(x) for h in self.heads], dim=-1)
+        out = self.dropout(self.proj(out))
+        return out
+
+class FeedForward(nn.Module):
+    def __init__(self, n_embd):
+        super().__init__()
+        self.net = nn.Sequential(
+            nn.Linear(n_embd, 4 * n_embd),
+            nn.ReLU(),
+            nn.Linear(4 * n_embd, n_embd),
+            nn.Dropout(dropout),
+        )
+    def forward(self, x):
+        return self.net(x)
+
+class Block(nn.Module):
+    def __init__(self, n_embd, n_head):
+        super().__init__()
+        head_size = n_embd // n_head
+        self.sa   = MultiHeadAttention(n_head, head_size)
+        self.ffwd = FeedForward(n_embd)
+        self.ln1  = nn.LayerNorm(n_embd)
+        self.ln2  = nn.LayerNorm(n_embd)
+
+    def forward(self, x):
+        x = x + self.sa(self.ln1(x))
+        x = x + self.ffwd(self.ln2(x))
+        return x
+
+class LanguageModel(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.token_embedding_table    = nn.Embedding(vocab_size, n_embd)
+        self.position_embedding_table = nn.Embedding(block_size, n_embd)
+        self.blocks = nn.Sequential(*[Block(n_embd, n_head=n_head) for _ in range(n_layer)])
+        self.ln_f   = nn.LayerNorm(n_embd)
+        self.lm_head = nn.Linear(n_embd, vocab_size)
+
+    def forward(self, idx, targets=None):
+        B, T = idx.shape
+        tok_emb = self.token_embedding_table(idx)
+        pos_emb = self.position_embedding_table(torch.arange(T, device=device))
+        x = tok_emb + pos_emb
+        x = self.blocks(x)
+        x = self.ln_f(x)
+        logits = self.lm_head(x)
+        return logits, None
+
+    def generate(self, idx, max_new_tokens):
+        for _ in range(max_new_tokens):
+            idx_cond = idx[:, -block_size:]
+            logits, _ = self(idx_cond)
+            logits = logits[:, -1, :]
+            probs = F.softmax(logits, dim=-1)
+            idx_next = torch.multinomial(probs, num_samples=1)
+            idx = torch.cat((idx, idx_next), dim=1)
+        return idx
+
+# 2. LOAD TRAINED WEIGHTS
+try:
+    print("Loading model weights from 'tulsidas_model.pth'...")
+    model = LanguageModel()
+    model.load_state_dict(torch.load('tulsidas_model.pth', map_location=device))
+    model = model.to(device)
+    model.eval() # Set to evaluation mode
+    print("Model loaded successfully!")
+except FileNotFoundError:
+    print("Error: 'tulsidas_model.pth' not found. Please run train_llm.py first.")
+    sys.exit(1)
+
+# 3. INTERACTIVE CHAT INTERFACE
+def chat():
+    print("\n" + "="*50)
+    print("🙏 MANAS (Model for Awadhi Natural Autoregressive Sequences) 🙏")
+    print("You can type in English (e.g. 'shri ram') or Devanagari (e.g. 'श्री राम').")
+    print("Type 'quit' or 'exit' to stop.")
+    print("="*50 + "\n")
+    
+    while True:
+        try:
+            user_input = input("You: ")
+            if user_input.lower() in ['quit', 'exit']:
+                break
+            if not user_input.strip():
+                continue
+            
+            # Detect if user typed in English letters
+            if re.search(r'[a-zA-Z]', user_input):
+                # Transliterate ITRANS (English) -> Devanagari
+                hindi_input = transliterate(user_input, sanscript.ITRANS, sanscript.DEVANAGARI)
+                print(f"[Auto-translated]: {hindi_input}")
+            else:
+                hindi_input = user_input
+            
+            # Start with the user's prompt
+            context = torch.tensor([encode(hindi_input)], dtype=torch.long, device=device)
+            
+            print("\nMANAS:")
+            generated_tokens = model.generate(context, max_new_tokens=500)[0].tolist()
+            print(decode(generated_tokens))
+            print("\n" + "-"*50)
+            
+        except KeyboardInterrupt:
+            break
+
+if __name__ == '__main__':
+    chat()