Initial upload of custom Hindi LM v1

README.md

@@ -0,0 +1,190 @@
+
+---
+license: mit # Or choose another appropriate license: https://huggingface.co/docs/hub/repositories-licenses
+language: hi
+tags:
+- hindi
+- text-generation
+- causal-lm
+- custom-model
+# Add more specific tags if applicable, e.g., based on training data domain
+pipeline_tag: text-generation
+# Specify model size if known from config
+
+---
+
+# Hindi Causal Language Model (convaiinnovations/hindi-foundational-model-base)
+
+This repository contains a custom-trained Hindi Causal Language Model.
+
+## Model Description
+
+*   **Architecture:** Custom Transformer (12 layers, hidden=768, 16 heads, ffn=3072, act=swiglu, norm=rmsnorm) based on the `HindiCausalLM` class with modified attention mechanisms and Hindi-specific optimizations including multi-resolution attention to capture both character-level and word-level patterns, morphology-aware feed-forward layers, and script-mix processing for Hindi-English code-mixing.
+*   **Language:** Hindi (hi)
+*   **Training Data:** The model was trained on a diverse corpus of 2.7 million high-quality Hindi text samples from multiple sources including IITB Parallel Corpus (1.2M sentences), Samanantar (750K samples), Oscar Hindi (450K sentences), CC-100 Hindi (300K sentences), Hindi Wikipedia (150K articles), Hindi news articles (100K pieces), XNLI Hindi (50K premise-hypothesis pairs), IndicGLUE (30K samples), and Hindi literature (5K passages).
+*   **Tokenizer:** SentencePiece trained on Hindi text (`tokenizer.model`). Vocab Size: 16000
+*   **Model Details:** Trained for [Unknown] epochs [Add more training details like batch size, learning rate if desired].
+
+## How to Use
+
+**⚠️ Important:** This model uses custom Python classes (`HindiCausalLM`, `HindiCausalLMConfig`, `SentencePieceTokenizerWrapper`) which are **not** part of the standard Hugging Face `transformers` library. To use this model, you **must** have the Python files defining these classes (e.g., `hindi_language_model.py`, `hindi_embeddings.py`) available in your Python environment.
+
+```python
+import os
+import json
+import torch
+import numpy as np
+from huggingface_hub import hf_hub_download
+
+# --- ENSURE THESE CLASSES ARE DEFINED OR IMPORTED ---
+# You MUST have hindi_language_model.py and hindi_embeddings.py available
+try:
+    from hindi_language_model import HindiCausalLM, HindiCausalLMConfig
+    from hindi_embeddings import SentencePieceTokenizerWrapper
+    print("Custom classes imported successfully.")
+except ImportError:
+    print("ERROR: Cannot import custom classes.")
+    print("Please place hindi_language_model.py and hindi_embeddings.py in your working directory or Python path.")
+    # Define minimal dummy classes to potentially allow script execution, but loading will fail
+    class SentencePieceTokenizerWrapper: pass
+    class HindiCausalLMConfig: pass
+    class HindiCausalLM(torch.nn.Module): pass
+    # Exit if classes are truly needed
+    # exit()
+
+
+# --- Configuration ---
+repo_id = "convaiinnovations/hindi-foundational-model-base"
+# model_dir = "./downloaded_model" # Example download location
+# os.makedirs(model_dir, exist_ok=True)
+# Use current directory if preferred
+model_dir = "."
+
+
+# --- Download Files ---
+print(f"Downloading files for {repo_id} to '{os.path.abspath(model_dir)}'...")
+try:
+    config_path = hf_hub_download(repo_id=repo_id, filename="config.json", local_dir=model_dir, local_dir_use_symlinks=False)
+    tokenizer_path = hf_hub_download(repo_id=repo_id, filename="tokenizer.model", local_dir=model_dir, local_dir_use_symlinks=False)
+    # Try safetensors first, then bin
+    using_safetensors = True
+    try: weights_path = hf_hub_download(repo_id=repo_id, filename="model.safetensors", local_dir=model_dir, local_dir_use_symlinks=False)
+    except Exception: # More specific: from huggingface_hub.utils import EntryNotFoundError
+        try: weights_path = hf_hub_download(repo_id=repo_id, filename="pytorch_model.bin", local_dir=model_dir, local_dir_use_symlinks=False); using_safetensors = False
+        except Exception as e_inner: raise FileNotFoundError(f"Could not download weights (.safetensors or .bin): {e_inner}") from e_inner
+except Exception as e: raise RuntimeError(f"Failed to download files from Hub: {e}") from e
+print("Files downloaded.")
+
+# --- Load Components ---
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+print(f"Using device: {device}")
+
+try:
+    # 1. Load Tokenizer
+    tokenizer = SentencePieceTokenizerWrapper(tokenizer_path) # Assumes constructor takes path
+
+    # 2. Load Config
+    with open(config_path, 'r', encoding='utf-8') as f: config_dict = json.load(f)
+    if hasattr(HindiCausalLMConfig, 'from_dict'): config = HindiCausalLMConfig.from_dict(config_dict)
+    else: config = HindiCausalLMConfig(**config_dict) # Assumes __init__ takes kwargs
+
+    # 3. Load Model
+    model = HindiCausalLM(config) # Instantiate model
+    if using_safetensors:
+        from safetensors.torch import load_file
+        state_dict = load_file(weights_path, device="cpu")
+    else:
+        state_dict = torch.load(weights_path, map_location="cpu")
+    model.load_state_dict(state_dict, strict=True) # Use strict=True after training
+    del state_dict
+    model.to(device)
+    model.eval()
+    print("Model and tokenizer loaded successfully.")
+
+except Exception as e:
+    print(f"ERROR: Failed loading model components: {e}")
+    # Add more specific error handling if needed
+    exit()
+
+
+# --- Example Inference ---
+prompt = "भारत की संस्कृति" # Example prompt
+max_new_tokens = 60 # Generate N new tokens
+temperature = 0.7
+top_k = 50
+seed = 42
+
+print(f"\nGenerating text for prompt: '{prompt}'...")
+
+torch.manual_seed(seed)
+np.random.seed(seed)
+if device.type == 'cuda': torch.cuda.manual_seed_all(seed)
+
+try:
+    # Encoding (use the correct method from your wrapper)
+    if hasattr(tokenizer, '__call__'): encoded = tokenizer(prompt, return_tensors=None); input_ids = encoded.get('input_ids')
+    elif hasattr(tokenizer, 'sp_model') and hasattr(tokenizer.sp_model, 'EncodeAsIds'): input_ids = tokenizer.sp_model.EncodeAsIds(prompt)
+    else: raise AttributeError("Tokenizer lacks encoding method.")
+    assert input_ids, "Encoding failed"
+
+    bos_id = getattr(tokenizer, 'bos_token_id', 1)
+    if bos_id is not None: input_ids = [bos_id] + input_ids
+
+    input_tensor = torch.tensor([input_ids], dtype=torch.long, device=device)
+    generated_ids = input_tensor
+
+    with torch.no_grad():
+        for _ in range(max_new_tokens):
+            outputs = model(input_ids=generated_ids)
+            # Access logits
+            if isinstance(outputs, dict) and 'logits' in outputs: logits = outputs['logits']
+            elif hasattr(outputs, 'logits'): logits = outputs.logits
+            else: raise TypeError("Model output format error.")
+
+            next_token_logits = logits[:, -1, :]
+            # Sampling
+            if temperature > 0: scaled_logits = next_token_logits / temperature
+            else: scaled_logits = next_token_logits
+            if top_k > 0: kth_vals, _ = torch.topk(scaled_logits, k=top_k); scaled_logits[scaled_logits < kth_vals[:, -1].unsqueeze(-1)] = -float("Inf")
+            probs = torch.softmax(scaled_logits, dim=-1)
+            next_token_id = torch.multinomial(probs, num_samples=1)
+            generated_ids = torch.cat([generated_ids, next_token_id], dim=1)
+            # Check EOS
+            eos_id = getattr(tokenizer, 'eos_token_id', 2)
+            if eos_id is not None and next_token_id.item() == eos_id: break
+
+    # Decoding
+    output_ids = generated_ids[0].cpu().tolist()
+    # Remove special tokens
+    if bos_id and output_ids and output_ids[0] == bos_id: output_ids = output_ids[1:]
+    if eos_id and output_ids and output_ids[-1] == eos_id: output_ids = output_ids[:-1]
+
+    # Use appropriate decode method
+    if hasattr(tokenizer, 'sp_model') and hasattr(tokenizer.sp_model, 'DecodeIds'): generated_text = tokenizer.sp_model.DecodeIds(output_ids)
+    elif hasattr(tokenizer, 'decode'): generated_text = tokenizer.decode(output_ids)
+    else: raise AttributeError("Tokenizer lacks decoding method.")
+
+    print("\n--- Generated Text ---")
+    # Print prompt + generated text for context
+    print(prompt + generated_text)
+    print("----------------------")
+
+except Exception as e:
+    print(f"\nERROR during example inference: {e}")
+```
+
+## Limitations and Biases
+
+This model was trained on a diverse corpus of Hindi text from sources including IITB Parallel Corpus, Samanantar, Oscar Hindi, CC-100 Hindi, Hindi Wikipedia, news articles, XNLI Hindi, IndicGLUE, and Hindi literature. As such, it may reflect biases present in that data, including potential cultural, gender, or regional biases found in these source materials.
+
+The model's performance is limited by its architecture (12 layers, hidden=768, 16 heads, ffn=3072, act=swiglu, norm=rmsnorm) and the size of the training dataset.
+
+It may generate repetitive, nonsensical, or factually incorrect text.
+
+The model uses a weighted pooling strategy with sensitivity to Hindi's SOV structure, but may still struggle with complex semantic relationships in longer texts.
+
+As noted in the DeepRAG research paper, the model may have particular difficulties with cultural concepts that lack direct English translations, idiomatic expressions specific to Hindi, and formal/informal speech distinctions.
+
+Please use this model responsibly.
+
+Model trained using custom scripts.

config.json

@@ -0,0 +1,23 @@
+{
+  "vocab_size": 16000,
+  "hidden_size": 768,
+  "num_hidden_layers": 12,
+  "num_attention_heads": 16,
+  "intermediate_size": 3072,
+  "hidden_dropout_prob": 0.1,
+  "attention_probs_dropout_prob": 0.1,
+  "max_position_embeddings": 512,
+  "layer_norm_eps": 1e-12,
+  "pad_token_id": 0,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "tie_word_embeddings": true,
+  "activation_function": "swiglu",
+  "normalization_layer": "rmsnorm",
+  "positional_encoding_type": "rope",
+  "unk_token_id": 3,
+  "architectures": [
+    "HindiCausalLM"
+  ],
+  "model_type": "hindi_causal_lm"
+}

hindi_embeddings.py

@@ -0,0 +1,730 @@
+import os
+import torch
+import json
+import numpy as np
+from torch import nn
+from torch.nn import functional as F
+import sentencepiece as spm
+from sklearn.metrics.pairwise import cosine_similarity
+from tqdm import tqdm
+import matplotlib.pyplot as plt
+from sklearn.manifold import TSNE
+
+# Tokenizer wrapper class
+class SentencePieceTokenizerWrapper:
+    def __init__(self, sp_model_path):
+        self.sp_model = spm.SentencePieceProcessor()
+        self.sp_model.Load(sp_model_path)
+        self.vocab_size = self.sp_model.GetPieceSize()
+        
+        # Special token IDs from tokenizer training
+        self.pad_token_id = 0
+        self.bos_token_id = 1
+        self.eos_token_id = 2
+        self.unk_token_id = 3
+        
+        # Set special tokens
+        self.pad_token = "<pad>"
+        self.bos_token = "<s>"
+        self.eos_token = "</s>"
+        self.unk_token = "<unk>"
+        self.mask_token = "<mask>"
+    
+    def __call__(self, text, padding=False, truncation=False, max_length=None, return_tensors=None):
+        # Handle both string and list inputs
+        if isinstance(text, str):
+            # Encode a single string
+            ids = self.sp_model.EncodeAsIds(text)
+            
+            # Handle truncation
+            if truncation and max_length and len(ids) > max_length:
+                ids = ids[:max_length]
+                
+            attention_mask = [1] * len(ids)
+            
+            # Handle padding
+            if padding and max_length:
+                padding_length = max(0, max_length - len(ids))
+                ids = ids + [self.pad_token_id] * padding_length
+                attention_mask = attention_mask + [0] * padding_length
+            
+            result = {
+                'input_ids': ids,
+                'attention_mask': attention_mask
+            }
+            
+            # Convert to tensors if requested
+            if return_tensors == 'pt':
+                import torch
+                result = {k: torch.tensor([v]) for k, v in result.items()}
+            
+            return result
+        
+        # Process a batch of texts
+        batch_encoded = [self.sp_model.EncodeAsIds(t) for t in text]
+        
+        # Apply truncation if needed
+        if truncation and max_length:
+            batch_encoded = [ids[:max_length] for ids in batch_encoded]
+        
+        # Create attention masks
+        batch_attention_mask = [[1] * len(ids) for ids in batch_encoded]
+        
+        # Apply padding if needed
+        if padding:
+            if max_length:
+                max_len = max_length
+            else:
+                max_len = max(len(ids) for ids in batch_encoded)
+            
+            # Pad sequences to max_len
+            batch_encoded = [ids + [self.pad_token_id] * (max_len - len(ids)) for ids in batch_encoded]
+            batch_attention_mask = [mask + [0] * (max_len - len(mask)) for mask in batch_attention_mask]
+        
+        result = {
+            'input_ids': batch_encoded,
+            'attention_mask': batch_attention_mask
+        }
+        
+        # Convert to tensors if requested
+        if return_tensors == 'pt':
+            import torch
+            result = {k: torch.tensor(v) for k, v in result.items()}
+        
+        return result
+
+# Model architecture components
+class MultiHeadAttention(nn.Module):
+    """Multi-headed attention mechanism"""
+    def __init__(self, config):
+        super().__init__()
+        self.num_attention_heads = config["num_attention_heads"]
+        self.attention_head_size = config["hidden_size"] // config["num_attention_heads"]
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        
+        # Query, Key, Value projections
+        self.query = nn.Linear(config["hidden_size"], self.all_head_size)
+        self.key = nn.Linear(config["hidden_size"], self.all_head_size)
+        self.value = nn.Linear(config["hidden_size"], self.all_head_size)
+        
+        # Output projection
+        self.output = nn.Sequential(
+            nn.Linear(self.all_head_size, config["hidden_size"]),
+            nn.Dropout(config["attention_probs_dropout_prob"])
+        )
+        
+        # Simplified relative position bias
+        self.max_position_embeddings = config["max_position_embeddings"]
+        self.relative_attention_bias = nn.Embedding(
+            2 * config["max_position_embeddings"] - 1, 
+            config["num_attention_heads"]
+        )
+        
+    def transpose_for_scores(self, x):
+        new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_shape)
+        return x.permute(0, 2, 1, 3)
+    
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_length = hidden_states.size()[:2]
+        
+        # Project inputs to queries, keys, and values
+        query_layer = self.transpose_for_scores(self.query(hidden_states))
+        key_layer = self.transpose_for_scores(self.key(hidden_states))
+        value_layer = self.transpose_for_scores(self.value(hidden_states))
+        
+        # Take the dot product between query and key to get the raw attention scores
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        
+        # Generate relative position matrix
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=hidden_states.device)
+        relative_position = position_ids.unsqueeze(1) - position_ids.unsqueeze(0)  # [seq_len, seq_len]
+        # Shift values to be >= 0
+        relative_position = relative_position + self.max_position_embeddings - 1  
+        # Ensure indices are within bounds
+        relative_position = torch.clamp(relative_position, 0, 2 * self.max_position_embeddings - 2)  
+        
+        # Get relative position embeddings [seq_len, seq_len, num_heads]
+        rel_attn_bias = self.relative_attention_bias(relative_position)  # [seq_len, seq_len, num_heads]
+        
+        # Reshape to add to attention heads [1, num_heads, seq_len, seq_len]
+        rel_attn_bias = rel_attn_bias.permute(2, 0, 1).unsqueeze(0)
+        
+        # Add to attention scores - now dimensions will match
+        attention_scores = attention_scores + rel_attn_bias
+        
+        # Scale attention scores
+        attention_scores = attention_scores / (self.attention_head_size ** 0.5)
+        
+        # Apply attention mask
+        if attention_mask is not None:
+            attention_scores = attention_scores + attention_mask
+        
+        # Normalize the attention scores to probabilities
+        attention_probs = F.softmax(attention_scores, dim=-1)
+        
+        # Apply dropout
+        attention_probs = F.dropout(attention_probs, p=0.1, training=self.training)
+        
+        # Apply attention to values
+        context_layer = torch.matmul(attention_probs, value_layer)
+        
+        # Reshape back to [batch_size, seq_length, hidden_size]
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_shape)
+        
+        # Final output projection
+        output = self.output(context_layer)
+        
+        return output
+
+class EnhancedTransformerLayer(nn.Module):
+    """Advanced transformer layer with pre-layer norm and enhanced attention"""
+    def __init__(self, config):
+        super().__init__()
+        self.attention_pre_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        self.attention = MultiHeadAttention(config)
+        
+        self.ffn_pre_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(config["hidden_size"], config["intermediate_size"]),
+            nn.GELU(),
+            nn.Dropout(config["hidden_dropout_prob"]),
+            nn.Linear(config["intermediate_size"], config["hidden_size"]),
+            nn.Dropout(config["hidden_dropout_prob"])
+        )
+        
+    def forward(self, hidden_states, attention_mask=None):
+        # Pre-layer norm for attention
+        attn_norm_hidden = self.attention_pre_norm(hidden_states)
+        
+        # Self-attention
+        attention_output = self.attention(attn_norm_hidden, attention_mask)
+        
+        # Residual connection
+        hidden_states = hidden_states + attention_output
+        
+        # Pre-layer norm for feed-forward
+        ffn_norm_hidden = self.ffn_pre_norm(hidden_states)
+        
+        # Feed-forward
+        ffn_output = self.ffn(ffn_norm_hidden)
+        
+        # Residual connection
+        hidden_states = hidden_states + ffn_output
+        
+        return hidden_states
+
+class AdvancedTransformerModel(nn.Module):
+    """Advanced Transformer model for inference"""
+    
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Embeddings
+        self.word_embeddings = nn.Embedding(
+            config["vocab_size"], 
+            config["hidden_size"], 
+            padding_idx=config["pad_token_id"]
+        )
+        
+        # Position embeddings
+        self.position_embeddings = nn.Embedding(config["max_position_embeddings"], config["hidden_size"])
+        
+        # Embedding dropout
+        self.embedding_dropout = nn.Dropout(config["hidden_dropout_prob"])
+        
+        # Transformer layers
+        self.layers = nn.ModuleList([
+            EnhancedTransformerLayer(config) for _ in range(config["num_hidden_layers"])
+        ])
+        
+        # Final layer norm
+        self.final_layer_norm = nn.LayerNorm(config["hidden_size"], eps=config["layer_norm_eps"])
+        
+    def forward(self, input_ids, attention_mask=None):
+        input_shape = input_ids.size()
+        batch_size, seq_length = input_shape
+        
+        # Get position ids
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        
+        # Get embeddings
+        word_embeds = self.word_embeddings(input_ids)
+        position_embeds = self.position_embeddings(position_ids)
+        
+        # Sum embeddings
+        embeddings = word_embeds + position_embeds
+        
+        # Apply dropout
+        embeddings = self.embedding_dropout(embeddings)
+        
+        # Default attention mask
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=input_ids.device)
+        
+        # Extended attention mask for transformer layers (1 for tokens to attend to, 0 for masked tokens)
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+        
+        # Apply transformer layers
+        hidden_states = embeddings
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, extended_attention_mask)
+        
+        # Final layer norm
+        hidden_states = self.final_layer_norm(hidden_states)
+        
+        return hidden_states
+
+class AdvancedPooling(nn.Module):
+    """Advanced pooling module supporting multiple pooling strategies"""
+    def __init__(self, config):
+        super().__init__()
+        self.pooling_mode = config["pooling_mode"]  # 'mean', 'max', 'cls', 'attention'
+        self.hidden_size = config["hidden_size"]
+        
+        # For attention pooling
+        if self.pooling_mode == 'attention':
+            self.attention_weights = nn.Linear(config["hidden_size"], 1)
+            
+        # For weighted pooling
+        elif self.pooling_mode == 'weighted':
+            self.weight_layer = nn.Linear(config["hidden_size"], 1)
+            
+    def forward(self, token_embeddings, attention_mask=None):
+        if attention_mask is None:
+            attention_mask = torch.ones_like(token_embeddings[:, :, 0])
+            
+        mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).float()
+        
+        if self.pooling_mode == 'cls':
+            # Use [CLS] token (first token)
+            pooled = token_embeddings[:, 0]
+            
+        elif self.pooling_mode == 'max':
+            # Max pooling
+            token_embeddings = token_embeddings.clone()
+            # Set padding tokens to large negative value to exclude them from max
+            token_embeddings[mask_expanded == 0] = -1e9
+            pooled = torch.max(token_embeddings, dim=1)[0]
+            
+        elif self.pooling_mode == 'attention':
+            # Attention pooling
+            weights = self.attention_weights(token_embeddings).squeeze(-1)
+            # Mask out padding tokens
+            weights = weights.masked_fill(attention_mask == 0, -1e9)
+            weights = F.softmax(weights, dim=1).unsqueeze(-1)
+            pooled = torch.sum(token_embeddings * weights, dim=1)
+            
+        elif self.pooling_mode == 'weighted':
+            # Weighted average pooling
+            weights = torch.sigmoid(self.weight_layer(token_embeddings)).squeeze(-1)
+            # Apply mask
+            weights = weights * attention_mask
+            # Normalize weights
+            sum_weights = torch.sum(weights, dim=1, keepdim=True)
+            sum_weights = torch.clamp(sum_weights, min=1e-9)
+            weights = weights / sum_weights
+            # Apply weights
+            pooled = torch.sum(token_embeddings * weights.unsqueeze(-1), dim=1)
+            
+        else:  # Default to mean pooling
+            # Mean pooling
+            sum_embeddings = torch.sum(token_embeddings * mask_expanded, dim=1)
+            sum_mask = torch.clamp(mask_expanded.sum(1), min=1e-9)
+            pooled = sum_embeddings / sum_mask
+            
+        # L2 normalize
+        pooled = F.normalize(pooled, p=2, dim=1)
+        
+        return pooled
+
+class SentenceEmbeddingModel(nn.Module):
+    """Complete sentence embedding model for inference"""
+    def __init__(self, config):
+        super(SentenceEmbeddingModel, self).__init__()
+        self.config = config
+        
+        # Create transformer model
+        self.transformer = AdvancedTransformerModel(config)
+        
+        # Create pooling module
+        self.pooling = AdvancedPooling(config)
+        
+        # Build projection module if needed
+        if "projection_dim" in config and config["projection_dim"] > 0:
+            self.use_projection = True
+            self.projection = nn.Sequential(
+                nn.Linear(config["hidden_size"], config["hidden_size"]),
+                nn.GELU(),
+                nn.Linear(config["hidden_size"], config["projection_dim"]),
+                nn.LayerNorm(config["projection_dim"], eps=config["layer_norm_eps"])
+            )
+        else:
+            self.use_projection = False
+            
+    def forward(self, input_ids, attention_mask=None):
+        # Get token embeddings from transformer
+        token_embeddings = self.transformer(input_ids, attention_mask)
+        
+        # Pool token embeddings
+        pooled_output = self.pooling(token_embeddings, attention_mask)
+        
+        # Apply projection if enabled
+        if self.use_projection:
+            pooled_output = self.projection(pooled_output)
+            pooled_output = F.normalize(pooled_output, p=2, dim=1)
+        
+        return pooled_output
+
+class HindiEmbedder:
+    def __init__(self, model_path="/home/ubuntu/output/hindi-embeddings-custom-tokenizer/final"):
+        """
+        Initialize the Hindi sentence embedder.
+        
+        Args:
+            model_path: Path to the model directory
+        """
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        print(f"Using device: {self.device}")
+        
+        # Load tokenizer - look for it in the model directory
+        tokenizer_path = os.path.join(model_path, "tokenizer.model")
+        
+        if not os.path.exists(tokenizer_path):
+            raise FileNotFoundError(f"Could not find tokenizer at {tokenizer_path}")
+            
+        self.tokenizer = SentencePieceTokenizerWrapper(tokenizer_path)
+        print(f"Loaded tokenizer from {tokenizer_path} with vocabulary size: {self.tokenizer.vocab_size}")
+        
+        # Load model config
+        config_path = os.path.join(model_path, "config.json")
+        with open(config_path, "r") as f:
+            self.config = json.load(f)
+        print(f"Loaded model config with hidden_size={self.config['hidden_size']}")
+        
+        # Load model
+        model_pt_path = os.path.join(model_path, "embedding_model.pt")
+        
+        try:
+            # Support both PyTorch 2.6+ and older versions
+            try:
+                checkpoint = torch.load(model_pt_path, map_location=self.device, weights_only=False)
+                print("Loaded model using PyTorch 2.6+ style loading")
+            except TypeError:
+                checkpoint = torch.load(model_pt_path, map_location=self.device)
+                print("Loaded model using older PyTorch style loading")
+                
+            # Create model
+            self.model = SentenceEmbeddingModel(self.config)
+            
+            # Load state dict
+            if "model_state_dict" in checkpoint:
+                state_dict = checkpoint["model_state_dict"]
+            else:
+                state_dict = checkpoint
+                
+            missing_keys, unexpected_keys = self.model.load_state_dict(state_dict, strict=False)
+            print(f"Loaded model with {len(missing_keys)} missing keys and {len(unexpected_keys)} unexpected keys")
+            
+            # Move to device
+            self.model.to(self.device)
+            self.model.eval()
+            print("Model loaded successfully and placed in evaluation mode")
+            
+        except Exception as e:
+            print(f"Error loading model: {e}")
+            raise RuntimeError(f"Failed to load the model: {e}")
+    
+    def encode(self, sentences, batch_size=32, normalize=True):
+        """
+        Encode sentences to embeddings.
+        
+        Args:
+            sentences: A string or list of strings to encode
+            batch_size: Batch size for encoding
+            normalize: Whether to normalize the embeddings
+            
+        Returns:
+            Numpy array of embeddings
+        """
+        # Handle single sentence
+        if isinstance(sentences, str):
+            sentences = [sentences]
+            
+        all_embeddings = []
+        
+        # Process in batches
+        with torch.no_grad():
+            for i in range(0, len(sentences), batch_size):
+                batch = sentences[i:i+batch_size]
+                
+                # Tokenize
+                inputs = self.tokenizer(
+                    batch,
+                    padding=True,
+                    truncation=True,
+                    max_length=self.config.get("max_position_embeddings", 128),
+                    return_tensors="pt"
+                )
+                
+                # Move to device
+                input_ids = inputs["input_ids"].to(self.device)
+                attention_mask = inputs["attention_mask"].to(self.device)
+                
+                # Get embeddings
+                embeddings = self.model(input_ids, attention_mask)
+                
+                # Move to CPU and convert to numpy
+                all_embeddings.append(embeddings.cpu().numpy())
+                
+        # Concatenate all embeddings
+        all_embeddings = np.vstack(all_embeddings)
+        
+        # Normalize if requested
+        if normalize:
+            all_embeddings = all_embeddings / np.linalg.norm(all_embeddings, axis=1, keepdims=True)
+            
+        return all_embeddings
+    
+    def compute_similarity(self, texts1, texts2=None):
+        """
+        Compute cosine similarity between texts.
+        
+        Args:
+            texts1: First set of texts
+            texts2: Second set of texts. If None, compute similarity matrix within texts1.
+            
+        Returns:
+            Similarity scores
+        """
+        # Convert single strings to lists for consistent handling
+        if isinstance(texts1, str):
+            texts1 = [texts1]
+        
+        if texts2 is not None and isinstance(texts2, str):
+            texts2 = [texts2]
+            
+        embeddings1 = self.encode(texts1)
+        
+        if texts2 is None:
+            # Compute similarity matrix within texts1
+            similarities = cosine_similarity(embeddings1)
+            return similarities
+        else:
+            # Compute similarity between texts1 and texts2
+            embeddings2 = self.encode(texts2)
+            
+            if len(texts1) == len(texts2):
+                # Compute pairwise similarity when the number of texts match
+                similarities = np.array([
+                    cosine_similarity([e1], [e2])[0][0]
+                    for e1, e2 in zip(embeddings1, embeddings2)
+                ])
+                
+                # If there's just one pair, return a scalar
+                if len(similarities) == 1:
+                    return similarities[0]
+                return similarities
+            else:
+                # Return full similarity matrix
+                return cosine_similarity(embeddings1, embeddings2)
+    
+    def search(self, query, documents, top_k=5):
+        """
+        Search for similar documents to a query.
+        
+        Args:
+            query: The query text
+            documents: List of documents to search
+            top_k: Number of top results to return
+            
+        Returns:
+            List of dictionaries with document and score
+        """
+        # Get embeddings
+        query_embedding = self.encode([query])[0]
+        document_embeddings = self.encode(documents)
+        
+        # Compute similarities
+        similarities = np.dot(document_embeddings, query_embedding)
+        
+        # Get top indices
+        top_indices = np.argsort(similarities)[-top_k:][::-1]
+        
+        # Return results
+        results = []
+        for idx in top_indices:
+            results.append({
+                "document": documents[idx],
+                "score": float(similarities[idx])
+            })
+            
+        return results
+        
+    def evaluate_similarity_samples(self):
+        """Evaluate model on some standard similarity examples for Hindi"""
+        test_pairs = [
+            (
+                "मुझे हिंदी में पढ़ना बहुत पसंद है।",
+                "मैं हिंदी किताबें बहुत पसंद करता हूँ।"
+            ),
+            (
+                "आज मौसम बहुत अच्छा है।",
+                "आज बारिश हो रही है।"
+            ),
+            (
+                "भारत एक विशाल देश है।",
+                "भारत में कई भाषाएँ बोली जाती हैं।"
+            ),
+            (
+                "कंप्यूटर विज्ञान एक रोचक विषय है।",
+                "मैं कंप्यूटर साइंस का छात्र हूँ।"
+            ),
+            (
+                "मैं रोज सुबह योग करता हूँ।",
+                "स्वस्थ रहने के लिए व्यायाम जरूरी है।"
+            ),
+            # Add contrasting pairs to test discrimination
+            (
+                "मुझे हिंदी में पढ़ना बहुत पसंद है।", 
+                "क्रिकेट भारत में सबसे लोकप्रिय खेल है।"
+            ),
+            (
+                "आज मौसम बहुत अच्छा है।",
+                "भारतीय व्यंजन दुनिया भर में मशहूर हैं।"
+            ),
+            (
+                "कंप्यूटर विज्ञान एक रोचक विषय है।",
+                "हिमालय दुनिया का सबसे ऊंचा पर्वत है।"
+            )
+        ]
+        
+        print("Evaluating model on standard similarity samples:")
+        for i, (text1, text2) in enumerate(test_pairs):
+            similarity = self.compute_similarity([text1], [text2])[0]
+            print(f"\nPair {i+1}:")
+            print(f"  Sentence 1: {text1}")
+            print(f"  Sentence 2: {text2}")
+            print(f"  Similarity: {similarity:.4f}")
+            
+        return
+        
+    def visualize_embeddings(self, sentences, labels=None, output_path="hindi_embeddings_visualization.png"):
+        """
+        Create a t-SNE visualization of the embeddings.
+        
+        Args:
+            sentences: List of sentences to visualize
+            labels: Optional list of labels for the points
+            output_path: Path to save the visualization
+            
+        Returns:
+            Path to the saved visualization
+        """
+        # Encode sentences
+        embeddings = self.encode(sentences)
+        
+        # Apply t-SNE
+        tsne = TSNE(n_components=2, random_state=42, perplexity=min(30, len(embeddings)-1))
+        reduced_embeddings = tsne.fit_transform(embeddings)
+        
+        # Create plot
+        plt.figure(figsize=(12, 10))
+        
+        # Plot points
+        scatter = plt.scatter(
+            reduced_embeddings[:, 0], 
+            reduced_embeddings[:, 1],
+            c=range(len(reduced_embeddings)),
+            cmap='viridis',
+            alpha=0.8,
+            s=100
+        )
+        
+        # Add labels if provided
+        if labels:
+            for i, label in enumerate(labels):
+                plt.annotate(
+                    label,
+                    (reduced_embeddings[i, 0], reduced_embeddings[i, 1]),
+                    fontsize=10,
+                    alpha=0.7
+                )
+        
+        plt.title("t-SNE Visualization of Hindi Sentence Embeddings", fontsize=16)
+        plt.xlabel("Dimension 1", fontsize=12)
+        plt.ylabel("Dimension 2", fontsize=12)
+        plt.colorbar(scatter, label="Sentence Index")
+        plt.grid(alpha=0.3)
+        
+        # Save the figure
+        plt.tight_layout()
+        plt.savefig(output_path, dpi=300, bbox_inches='tight')
+        plt.close()
+        
+        print(f"Visualization saved to {output_path}")
+        return output_path
+
+def main():
+    # Create embedder
+    embedder = HindiEmbedder()
+    
+    # Run sample evaluation
+    embedder.evaluate_similarity_samples()
+    
+    # Example of semantic search
+    print("\nSemantic Search Example:")
+    query = "भारत की संस्कृति"
+    documents = [
+        "भारतीय संस्कृति दुनिया की सबसे प्राचीन संस्कृतियों में से एक है।",
+        "भारत की आबादी 1.3 अरब से अधिक है।",
+        "हिमालय पर्वत श्रृंखला भारत के उत्तर में स्थित है।",
+        "भारतीय व्यंजन में मसालों का प्रयोग किया जाता है।",
+        "भारत में 22 आधिकारिक भाषाएँ हैं।",
+        "संस्कृति लोगों के रहन-सहन का तरीका है।",
+        "भारत के विभिन्न राज्यों की अपनी अलग संस्कृति है।",
+        "रामायण और महाभारत भारतीय संस्कृति के महत्वपूर्ण हिस्से हैं।",
+    ]
+    
+    results = embedder.search(query, documents)
+    
+    print(f"Query: {query}")
+    print("Top results:")
+    for i, result in enumerate(results):
+        print(f"{i+1}. Score: {result['score']:.4f}")
+        print(f"   {result['document']}")
+    
+    # Create visualization example
+    print("\nCreating embedding visualization...")
+    visualization_sentences = [
+        "मुझे हिंदी में पढ़ना बहुत पसंद है।",
+        "मैं हिंदी किताबें बहुत पसंद करता हूँ।",
+        "आज मौसम बहुत अच्छा है।",
+        "आज बारिश हो रही है।",
+        "भारत एक विशाल देश है।",
+        "भारत में कई भाषाएँ बोली जाती हैं।",
+        "कंप्यूटर विज्ञान एक रोचक विषय है।",
+        "मैं कंप्यूटर साइंस का छात्र हूँ।",
+        "क्रिकेट भारत में सबसे लोकप्रिय खेल है।",
+        "भारतीय व्यंजन दुनिया भर में मशहूर हैं।",
+        "हिमालय दुनिया का सबसे ऊंचा पर्वत है।",
+        "गंगा भारत की सबसे पवित्र नदी है।",
+        "दिल्ली भारत की राजधानी है।",
+        "मुंबई भारत का आर्थिक केंद्र है।",
+        "तमिल, तेलुगु, कन्नड़ और मलयालम दक्षिण भारत की प्रमुख भाषाएँ हैं।"
+    ]
+    
+    labels = ["पढ़ना", "किताबें", "मौसम", "बारिश", "भारत", "भाषाएँ", "क��प्यूटर", 
+              "छात्र", "क्रिकेट", "व्यंजन", "हिमालय", "गंगा", "दिल्ली", "मुंबई", "भाषाएँ"]
+    
+    embedder.visualize_embeddings(visualization_sentences, labels)
+
+if __name__ == "__main__":
+    main()

hindi_language_model.py

@@ -0,0 +1,547 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import math
+from typing import Optional, Tuple, List, Dict, Any, Union
+
+class HindiCausalLMConfig:
+    """Configuration class for Hindi Causal Language Model"""
+    def __init__(
+        self,
+        vocab_size: int = 32000,
+        hidden_size: int = 768,
+        num_hidden_layers: int = 12,
+        num_attention_heads: int = 12,
+        intermediate_size: int = 3072,
+        hidden_dropout_prob: float = 0.1,
+        attention_probs_dropout_prob: float = 0.1,
+        max_position_embeddings: int = 512,
+        layer_norm_eps: float = 1e-12,
+        pad_token_id: int = 0,
+        bos_token_id: int = 1,
+        eos_token_id: int = 2,
+        tie_word_embeddings: bool = True,
+        **kwargs
+    ):
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.layer_norm_eps = layer_norm_eps
+        self.pad_token_id = pad_token_id
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.tie_word_embeddings = tie_word_embeddings
+        
+        # Add any additional kwargs as attributes
+        for key, value in kwargs.items():
+            setattr(self, key, value)
+    
+    @classmethod
+    def from_embedding_config(cls, config_dict, **kwargs):
+        """Create LM config from embedding model config"""
+        # Check if override parameters are provided
+        override_params = {}
+        for key in ["num_hidden_layers", "hidden_size", "num_attention_heads", 
+                   "intermediate_size", "max_position_embeddings", "vocab_size"]:
+            if key in kwargs:
+                override_params[key] = kwargs.pop(key)
+        
+        # Get hidden size first to calculate appropriate number of attention heads
+        hidden_size = override_params.get("hidden_size", config_dict.get("hidden_size", 768))
+        
+        # If num_attention_heads is not provided, choose a value that divides hidden_size evenly
+        if "num_attention_heads" not in override_params:
+            # Default options to try: 12, 16, 8, 4
+            for heads in [12, 16, 8, 4]:
+                if hidden_size % heads == 0:
+                    print(f"Automatically setting num_attention_heads to {heads} to match hidden_size {hidden_size}")
+                    override_params["num_attention_heads"] = heads
+                    break
+            
+            # If none of the defaults work, find the largest factor <= 32
+            if "num_attention_heads" not in override_params:
+                # Find the largest factor of hidden_size that is <= 32
+                for heads in range(min(32, hidden_size), 0, -1):
+                    if hidden_size % heads == 0:
+                        print(f"Automatically setting num_attention_heads to {heads} to match hidden_size {hidden_size}")
+                        override_params["num_attention_heads"] = heads
+                        break
+        
+        # Build the config, with overrides taking precedence
+        config_params = {
+            "vocab_size": override_params.get("vocab_size", config_dict.get("vocab_size", 32000)),
+            "hidden_size": hidden_size,
+            "num_hidden_layers": override_params.get("num_hidden_layers", config_dict.get("num_hidden_layers", 12)),
+            "num_attention_heads": override_params.get("num_attention_heads", config_dict.get("num_attention_heads", 12)),
+            "intermediate_size": override_params.get("intermediate_size", config_dict.get("intermediate_size", 3072)),
+            "hidden_dropout_prob": config_dict.get("hidden_dropout_prob", 0.1),
+            "attention_probs_dropout_prob": config_dict.get("attention_probs_dropout_prob", 0.1),
+            "max_position_embeddings": override_params.get("max_position_embeddings", 
+                                                          config_dict.get("max_position_embeddings", 512)),
+            "layer_norm_eps": config_dict.get("layer_norm_eps", 1e-12),
+            "pad_token_id": config_dict.get("pad_token_id", 0),
+        }
+        
+        # Verify that hidden_size is divisible by num_attention_heads
+        if config_params["hidden_size"] % config_params["num_attention_heads"] != 0:
+            raise ValueError(
+                f"Hidden size ({config_params['hidden_size']}) must be divisible by the number of attention "
+                f"heads ({config_params['num_attention_heads']})"
+            )
+        
+        # Add remaining kwargs
+        config_params.update(kwargs)
+        
+        # Create and return the config
+        lm_config = cls(**config_params)
+        return lm_config
+    
+    def to_dict(self):
+        """Convert config to dictionary"""
+        return {k: v for k, v in self.__dict__.items()}
+
+class CausalSelfAttention(nn.Module):
+    """Causal self-attention layer"""
+    def __init__(self, config):
+        super().__init__()
+        assert config.hidden_size % config.num_attention_heads == 0
+        
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = config.hidden_size // config.num_attention_heads
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        
+        # Query, Key, Value projections
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+        
+        # Output projection
+        self.output = nn.Sequential(
+            nn.Linear(self.all_head_size, config.hidden_size),
+            nn.Dropout(config.attention_probs_dropout_prob)
+        )
+        
+        # Causal mask to ensure that attention is only applied to the left in the input sequence
+        self.register_buffer(
+            "causal_mask",
+            torch.triu(
+                torch.ones(config.max_position_embeddings, config.max_position_embeddings) * -1e10, 
+                diagonal=1
+            )
+        )
+        
+    def transpose_for_scores(self, x):
+        # Reshape from [batch_size, seq_length, hidden_size] to [batch_size, seq_length, num_heads, head_size]
+        new_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_shape)
+        # Transpose to [batch_size, num_heads, seq_length, head_size]
+        return x.permute(0, 2, 1, 3)
+    
+    def forward(self, hidden_states, attention_mask=None):
+        batch_size, seq_length = hidden_states.size()[:2]
+        
+        # Project inputs to queries, keys, and values
+        query_layer = self.transpose_for_scores(self.query(hidden_states))
+        key_layer = self.transpose_for_scores(self.key(hidden_states))
+        value_layer = self.transpose_for_scores(self.value(hidden_states))
+        
+        # Scale dot-product attention
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        
+        # Apply causal mask - prevents attending to future tokens
+        causal_mask = self.causal_mask[:seq_length, :seq_length]
+        attention_scores = attention_scores + causal_mask
+        
+        # Apply attention mask if provided
+        if attention_mask is not None:
+            # Expand mask to match attention_scores shape
+            attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+            attention_mask = (1.0 - attention_mask) * -10000.0
+            attention_scores = attention_scores + attention_mask
+        
+        # Softmax normalization
+        attention_probs = F.softmax(attention_scores, dim=-1)
+        attention_probs = F.dropout(attention_probs, p=0.1, training=self.training)
+        
+        # Apply attention to values
+        context_layer = torch.matmul(attention_probs, value_layer)
+        
+        # Reshape back to [batch_size, seq_length, hidden_size]
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_shape)
+        
+        # Final output projection
+        output = self.output(context_layer)
+        
+        return output
+
+class TransformerBlock(nn.Module):
+    """Transformer block with causal attention for language modeling"""
+    def __init__(self, config):
+        super().__init__()
+        self.attention = CausalSelfAttention(config)
+        self.attention_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        
+        # Feed-forward network
+        self.ffn = nn.Sequential(
+            nn.Linear(config.hidden_size, config.intermediate_size),
+            nn.GELU(),
+            nn.Linear(config.intermediate_size, config.hidden_size),
+            nn.Dropout(config.hidden_dropout_prob)
+        )
+        self.ffn_layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        
+    def forward(self, hidden_states, attention_mask=None):
+        # Self-attention block with residual connection and layer norm
+        attn_output = self.attention(hidden_states, attention_mask)
+        hidden_states = self.attention_layernorm(hidden_states + attn_output)
+        
+        # Feed-forward block with residual connection and layer norm
+        ffn_output = self.ffn(hidden_states)
+        hidden_states = self.ffn_layernorm(hidden_states + ffn_output)
+        
+        return hidden_states
+
+class HindiCausalLM(nn.Module):
+    """Hindi Causal Language Model for text generation"""
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        
+        # Embeddings
+        self.token_embeddings = nn.Embedding(
+            config.vocab_size, 
+            config.hidden_size, 
+            padding_idx=config.pad_token_id
+        )
+        self.position_embeddings = nn.Embedding(
+            config.max_position_embeddings, 
+            config.hidden_size
+        )
+        self.embedding_dropout = nn.Dropout(config.hidden_dropout_prob)
+        
+        # Transformer layers
+        self.layers = nn.ModuleList([
+            TransformerBlock(config) for _ in range(config.num_hidden_layers)
+        ])
+        
+        # LM head
+        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        
+        # Tie weights if configured
+        if config.tie_word_embeddings:
+            self.lm_head.weight = self.token_embeddings.weight
+        
+        # Initialize weights
+        self.apply(self._init_weights)
+    
+    def _init_weights(self, module):
+        """Initialize weights with small random values"""
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if isinstance(module, nn.Linear) and module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+    
+    def get_input_embeddings(self):
+        return self.token_embeddings
+    
+    def set_input_embeddings(self, new_embeddings):
+        self.token_embeddings = new_embeddings
+    
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        labels=None,
+        return_dict=True
+    ):
+        device = input_ids.device
+        batch_size, seq_length = input_ids.size()
+        
+        # Create position ids
+        position_ids = torch.arange(seq_length, dtype=torch.long, device=device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+        
+        # Get embeddings
+        inputs_embeds = self.token_embeddings(input_ids)
+        position_embeds = self.position_embeddings(position_ids)
+        
+        # Sum token and position embeddings
+        hidden_states = inputs_embeds + position_embeds
+        hidden_states = self.embedding_dropout(hidden_states)
+        
+        # Default attention mask (all tokens can be attended to)
+        if attention_mask is None:
+            attention_mask = torch.ones(batch_size, seq_length, device=device)
+        
+        # Apply transformer layers
+        for layer in self.layers:
+            hidden_states = layer(hidden_states, attention_mask)
+        
+        # Language model head
+        lm_logits = self.lm_head(hidden_states)
+        
+        loss = None
+        if labels is not None:
+            # Move labels to correct device
+            labels = labels.to(device)
+            
+            # Shift so that tokens < n predict n
+            shift_logits = lm_logits[..., :-1, :].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
+            
+            # Flatten the tokens
+            loss_fct = nn.CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(-1, shift_logits.size(-1)),
+                shift_labels.view(-1)
+            )
+        
+        if return_dict:
+            return {
+                "logits": lm_logits,
+                "loss": loss,
+                "hidden_states": hidden_states
+            }
+        
+        return (lm_logits, loss, hidden_states)
+    
+    def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **kwargs):
+        """Prepare inputs for text generation"""
+        # Only keep inputs needed for forward pass
+        inputs = {
+            "input_ids": input_ids,
+        }
+        
+        # Add attention mask if provided
+        if attention_mask is not None:
+            inputs["attention_mask"] = attention_mask
+            
+        return inputs
+    
+    @staticmethod
+    def _reorder_cache(past, beam_idx):
+        """Reorder cached hidden states for beam search generation"""
+        reordered_past = []
+        for layer_past in past:
+            reordered_past.append(
+                tuple(past_state.index_select(0, beam_idx) for past_state in layer_past)
+            )
+        return reordered_past
+    
+    def save_pretrained(self, save_directory):
+        """Save model and config to directory"""
+        import os
+        import json
+        import torch
+        
+        os.makedirs(save_directory, exist_ok=True)
+        
+        # Save config
+        config_path = os.path.join(save_directory, "config.json")
+        with open(config_path, "w", encoding="utf-8") as f:
+            json.dump(self.config.to_dict(), f, indent=2)
+        
+        # Save model weights
+        model_path = os.path.join(save_directory, "pytorch_model.bin")
+        torch.save(self.state_dict(), model_path)
+        
+        return [config_path, model_path]
+    
+    @classmethod
+    def from_pretrained(cls, model_path):
+        """Load model and config from directory"""
+        import os
+        import json
+        import torch
+        
+        # Load config
+        config_path = os.path.join(model_path, "config.json")
+        with open(config_path, "r", encoding="utf-8") as f:
+            config_dict = json.load(f)
+        
+        # Create config object
+        config = HindiCausalLMConfig(**config_dict)
+        
+        # Create model
+        model = cls(config)
+        
+        # Load model weights
+        model_path = os.path.join(model_path, "pytorch_model.bin")
+        model.load_state_dict(torch.load(model_path, map_location="cpu"))
+        
+        return model
+
+class HindiTextGenerator:
+    """Text generation utility for HindiCausalLM"""
+    def __init__(self, model, tokenizer):
+        self.model = model
+        self.tokenizer = tokenizer
+        self.device = next(model.parameters()).device
+    
+    def generate(
+        self,
+        prompt,
+        max_length=100,
+        temperature=1.0,
+        top_k=50,
+        top_p=0.95,
+        repetition_penalty=1.0,
+        do_sample=True,
+        num_return_sequences=1,
+        **kwargs
+    ):
+        """Generate text from a prompt"""
+        # Encode the prompt
+        input_ids = self.tokenizer(
+            prompt, 
+            return_tensors="pt",
+            truncation=True,
+            max_length=self.model.config.max_position_embeddings - max_length
+        )["input_ids"].to(self.device)
+        
+        # Set the model to evaluation mode
+        self.model.eval()
+        
+        # Set generation parameters
+        gen_kwargs = {
+            "max_length": input_ids.shape[1] + max_length,
+            "temperature": temperature,
+            "top_k": top_k,
+            "top_p": top_p,
+            "repetition_penalty": repetition_penalty,
+            "do_sample": do_sample,
+            "num_return_sequences": num_return_sequences,
+            **kwargs
+        }
+        
+        # Generate text
+        with torch.no_grad():
+            output_sequences = self._generate_text(input_ids, **gen_kwargs)
+        
+        # Decode generated sequences
+        generated_texts = []
+        for sequence in output_sequences:
+            # Remove the prompt from the generated text
+            sequence = sequence[input_ids.shape[1]:]
+            text = self.tokenizer.sp_model.DecodeIds(sequence.tolist())
+            generated_texts.append(text)
+        
+        if num_return_sequences == 1:
+            return generated_texts[0]
+        
+        return generated_texts
+    
+    def _generate_text(
+        self,
+        input_ids,
+        max_length,
+        temperature=1.0,
+        top_k=50,
+        top_p=0.95,
+        repetition_penalty=1.0,
+        do_sample=True,
+        num_return_sequences=1,
+        pad_token_id=None,
+        eos_token_id=None,
+        **kwargs
+    ):
+        """Core text generation logic"""
+        # Set pad_token_id and eos_token_id
+        pad_token_id = pad_token_id if pad_token_id is not None else self.model.config.pad_token_id
+        eos_token_id = eos_token_id if eos_token_id is not None else self.model.config.eos_token_id
+        
+        batch_size = input_ids.shape[0]
+        
+        # Create attention mask
+        attention_mask = torch.ones_like(input_ids)
+        
+        # Clone the input_ids for each return sequence
+        input_ids = input_ids.repeat(num_return_sequences, 1)
+        attention_mask = attention_mask.repeat(num_return_sequences, 1)
+        
+        unfinished_sequences = torch.ones(input_ids.shape[0], dtype=torch.long, device=input_ids.device)
+        
+        # Keep track of which sequences are already finished
+        cur_len = input_ids.shape[1]
+        
+        while cur_len < max_length:
+            # Prepare model inputs
+            model_inputs = self.model.prepare_inputs_for_generation(
+                input_ids, attention_mask=attention_mask
+            )
+            
+            # Forward pass
+            outputs = self.model(**model_inputs, return_dict=True)
+            next_token_logits = outputs["logits"][:, -1, :]
+            
+            # Apply temperature scaling
+            next_token_logits = next_token_logits / temperature
+            
+            # Apply repetition penalty
+            if repetition_penalty != 1.0:
+                for i in range(input_ids.shape[0]):
+                    for token_id in set(input_ids[i].tolist()):
+                        next_token_logits[i, token_id] /= repetition_penalty
+            
+            # Filter logits using top-k and top-p sampling
+            if do_sample:
+                # Top-k filtering
+                if top_k > 0:
+                    indices_to_remove = next_token_logits < torch.topk(next_token_logits, top_k)[0][..., -1, None]
+                    next_token_logits[indices_to_remove] = -float("Inf")
+                
+                # Top-p (nucleus) filtering
+                if top_p < 1.0:
+                    sorted_logits, sorted_indices = torch.sort(next_token_logits, descending=True)
+                    cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+                    
+                    # Remove tokens with cumulative probability above the threshold
+                    sorted_indices_to_remove = cumulative_probs > top_p
+                    # Shift the indices to the right to keep the first token above threshold
+                    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+                    sorted_indices_to_remove[..., 0] = 0
+                    
+                    for i in range(next_token_logits.shape[0]):
+                        indices_to_remove = sorted_indices[i][sorted_indices_to_remove[i]]
+                        next_token_logits[i, indices_to_remove] = -float("Inf")
+                
+                # Sample from the filtered distribution
+                probabilities = F.softmax(next_token_logits, dim=-1)
+                next_tokens = torch.multinomial(probabilities, 1).squeeze(1)
+            else:
+                # Greedy decoding
+                next_tokens = torch.argmax(next_token_logits, dim=-1)
+            
+            # Update unfinished sequences based on EOS token
+            if eos_token_id is not None:
+                # Set the unfinished flag to 0 if the sequence reaches EOS
+                unfinished_sequences = unfinished_sequences.mul(
+                    next_tokens.ne(eos_token_id).long()
+                )
+            
+            # Check if all sequences are finished
+            if unfinished_sequences.max() == 0:
+                break
+                
+            # Concatenate next tokens to input_ids
+            input_ids = torch.cat([input_ids, next_tokens.unsqueeze(-1)], dim=-1)
+            
+            # Expand attention mask
+            attention_mask = torch.cat(
+                [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+            )
+            
+            cur_len += 1
+            
+        return input_ids

model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ebf86827285c98325d2434abc44fd95f79fcc56161e3b7019faa49799e9eb7c
+size 452698216

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f1658727f4ea5c571f69a60f6defcd180014a1d69be6e9c1ec360d9510aa6b5e
+size 642200

tokenizer_config.json

@@ -0,0 +1,9 @@
+{
+  "tokenizer_class": "SentencePieceTokenizerWrapper",
+  "vocab_size": 16000,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 0,
+  "unk_token_id": 3,
+  "model_max_length": 512
+}