Upload modeling_pawan_embd.py with huggingface_hub

modeling_pawan_embd.py

@@ -0,0 +1,124 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from transformers import PreTrainedModel
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+class PawanEmbdModel(PreTrainedModel):
+    """
+    PawanEmbd Model - A lightweight embedding model for sentence embeddings.
+
+    This model outputs normalized embeddings suitable for semantic similarity tasks.
+    """
+
+    config_class = PawanEmbdConfig
+    base_model_prefix = "pawan_embd"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.config = config
+        self.hidden_size = config.hidden_size
+        self.output_size = config.output_size
+
+        # Token + Position embeddings
+        self.token_embedding = nn.Embedding(config.vocab_size, config.hidden_size)
+        self.position_embedding = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.dropout = nn.Dropout(config.dropout)
+        self.layer_norm = nn.LayerNorm(config.hidden_size)
+
+        # Transformer encoder
+        encoder_layer = nn.TransformerEncoderLayer(
+            d_model=config.hidden_size,
+            nhead=config.num_heads,
+            dim_feedforward=config.intermediate_size,
+            dropout=config.dropout,
+            activation='gelu',
+            batch_first=True,
+            norm_first=True
+        )
+        self.encoder = nn.TransformerEncoder(encoder_layer, num_layers=config.num_layers)
+
+        # Projection to output size
+        self.projection = nn.Sequential(
+            nn.Linear(config.hidden_size, config.hidden_size * 2),
+            nn.GELU(),
+            nn.Dropout(config.dropout),
+            nn.Linear(config.hidden_size * 2, config.output_size)
+        )
+
+        # Initialize weights
+        self.post_init()
+
+    def _init_weights(self, module):
+        """Initialize the weights"""
+        if isinstance(module, nn.Linear):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=0.02)
+        elif isinstance(module, nn.LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        attention_mask: torch.Tensor = None,
+        return_dict: bool = True,
+        normalize: bool = True
+    ):
+        """
+        Args:
+            input_ids: [batch_size, seq_len]
+            attention_mask: [batch_size, seq_len]
+            return_dict: Whether to return a ModelOutput object
+            normalize: Whether to L2-normalize the embeddings
+
+        Returns:
+            If return_dict=True: BaseModelOutputWithPooling
+            If return_dict=False: tuple of (last_hidden_state, pooler_output)
+        """
+        batch_size, seq_len = input_ids.shape
+
+        # Generate position IDs
+        position_ids = torch.arange(seq_len, dtype=torch.long, device=input_ids.device)
+        position_ids = position_ids.unsqueeze(0).expand(batch_size, -1)
+
+        # Embeddings
+        token_embeds = self.token_embedding(input_ids)
+        position_embeds = self.position_embedding(position_ids)
+        embeddings = self.dropout(self.layer_norm(token_embeds + position_embeds))
+
+        # Attention mask for transformer (convert 1/0 to True/False)
+        if attention_mask is not None:
+            attention_mask = attention_mask == 0  # True = masked position
+
+        # Transformer encoding
+        encoded = self.encoder(embeddings, src_key_padding_mask=attention_mask)
+
+        # CLS pooling (take first token)
+        cls_output = encoded[:, 0]
+
+        # Project to output dimension
+        pooler_output = self.projection(cls_output)
+
+        # Normalize embeddings
+        if normalize:
+            pooler_output = F.normalize(pooler_output, p=2, dim=-1)
+
+        if not return_dict:
+            return (encoded, pooler_output)
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=encoded,
+            pooler_output=pooler_output,
+            hidden_states=None,
+            attentions=None
+        )
+
+    def count_parameters(self):
+        """Count trainable parameters"""
+        return sum(p.numel() for p in self.parameters() if p.requires_grad)
+