Update app.py

app.py

@@ -9,7 +9,7 @@ from huggingface_hub import HfApi, hf_hub_download
 # Configuration constants
 MODEL_REPO = "Gajendra5490/Scrached_Trained_Model"
 CURRENT_USER = "gajendra82"
-CURRENT_UTC = "2025-05-06 15:52:11"
+CURRENT_UTC = "2025-05-06 15:54:34"
 
 def setup_logging():
     logging.basicConfig(
@@ -23,6 +23,23 @@ def setup_logging():
 
 logger = setup_logging()
 
+class PositionalEncoding(torch.nn.Module):
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super().__init__()
+        self.dropout = torch.nn.Dropout(p=dropout)
+        
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-torch.log(torch.tensor(10000.0)) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+        
+    def forward(self, x):
+        x = x + self.pe[:, :x.size(1)]
+        return self.dropout(x)
+
 class ImprovedTransformer(torch.nn.Module):
     def __init__(
         self,
@@ -39,14 +56,7 @@ class ImprovedTransformer(torch.nn.Module):
         
         self.d_model = d_model
         self.embedding = torch.nn.Embedding(vocab_size, d_model)
-        
-        # Position encoding
-        position = torch.arange(max_seq_length).unsqueeze(1)
-        div_term = torch.exp(torch.arange(0, d_model, 2) * (-torch.log(torch.tensor(10000.0)) / d_model))
-        pe = torch.zeros(max_seq_length, 1, d_model)
-        pe[:, 0, 0::2] = torch.sin(position * div_term)
-        pe[:, 0, 1::2] = torch.cos(position * div_term)
-        self.register_buffer('pe', pe)
+        self.pos_encoder = PositionalEncoding(d_model, dropout)
         
         # Main transformer
         self.transformer = torch.nn.Transformer(
@@ -59,29 +69,32 @@ class ImprovedTransformer(torch.nn.Module):
             batch_first=True
         )
         
-        # Output layers
+        # Output layer
         self.output_layer = torch.nn.Linear(d_model, vocab_size)
         self.norm = torch.nn.LayerNorm(d_model)
         
     def forward(self, src, tgt):
-        # Create masks
-        src_mask = self.transformer.generate_square_subsequent_mask(src.size(1)).to(src.device)
+        # Create padding masks
+        src_key_padding_mask = (src == 0).to(src.device)
+        tgt_key_padding_mask = (tgt == 0).to(tgt.device)
+        
+        # Create causal mask for target
         tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(1)).to(tgt.device)
         
-        # Embeddings
-        src = self.embedding(src) * torch.sqrt(torch.tensor(self.d_model))
-        tgt = self.embedding(tgt) * torch.sqrt(torch.tensor(self.d_model))
+        # Embeddings and positional encoding
+        src = self.embedding(src) * torch.sqrt(torch.tensor(self.d_model, dtype=torch.float))
+        tgt = self.embedding(tgt) * torch.sqrt(torch.tensor(self.d_model, dtype=torch.float))
         
-        # Add positional encoding
-        src = src + self.pe[:src.size(1)].transpose(0, 1)
-        tgt = tgt + self.pe[:tgt.size(1)].transpose(0, 1)
+        src = self.pos_encoder(src)
+        tgt = self.pos_encoder(tgt)
         
         # Transform
         output = self.transformer(
             src,
             tgt,
-            src_mask=src_mask,
-            tgt_mask=tgt_mask
+            tgt_mask=tgt_mask,
+            src_key_padding_mask=src_key_padding_mask,
+            tgt_key_padding_mask=tgt_key_padding_mask
         )
         
         # Output processing
@@ -115,31 +128,32 @@ class ModelInference:
                 token=token
             )
             
-            # Load tokenizer first
+            # Load model data first to get configuration
+            self.logger.info("Loading model data...")
+            model_data = torch.load(
+                model_path,
+                map_location=self.device
+            )
+            
+            # Load tokenizer
             self.logger.info("Loading tokenizer...")
             with open(tokenizer_path, 'r', encoding='utf-8') as f:
                 tokenizer_data = json.load(f)
             
-            # Initialize tokenizer
+            # Initialize tokenizer with the same vocabulary size as the saved model
             self.vocab = tokenizer_data['vocab']
+            vocab_size = len(self.vocab)  # Use exact vocab size from saved model
+            
             self.special_tokens = {
-                "<user>": len(self.vocab),
-                "<assistant>": len(self.vocab) + 1,
-                "<sep>": len(self.vocab) + 2,
-                "<eos>": len(self.vocab) + 3
+                "<user>": vocab_size,
+                "<assistant>": vocab_size + 1,
+                "<sep>": vocab_size + 2,
+                "<eos>": vocab_size + 3
             }
             
-            # Load model
-            self.logger.info("Loading model...")
-            model_data = torch.load(
-                model_path,
-                map_location=self.device
-            )
-            
-            # Initialize model with correct vocab size
-            vocab_size = len(self.vocab) + len(self.special_tokens)
+            # Initialize model with exact vocab size
             self.model = ImprovedTransformer(
-                vocab_size=vocab_size,
+                vocab_size=vocab_size + 4,  # Add exactly 4 special tokens
                 d_model=512,
                 nhead=8,
                 num_encoder_layers=3,
@@ -147,8 +161,19 @@ class ModelInference:
                 dim_feedforward=2048
             ).to(self.device)
             
+            # Fix state dict keys
+            fixed_state_dict = {}
+            for k, v in model_data['model_state_dict'].items():
+                if k.startswith('transformer.'):
+                    fixed_state_dict[k] = v
+                elif k.startswith('pos_encoder.'):
+                    if k == 'pos_encoder.pe':
+                        fixed_state_dict['pos_encoder.pe'] = v
+                else:
+                    fixed_state_dict[k] = v
+            
             # Load state dict
-            self.model.load_state_dict(model_data['model_state_dict'])
+            self.model.load_state_dict(fixed_state_dict, strict=True)
             self.model.eval()
             
             self.logger.info("Model loaded successfully")
@@ -234,6 +259,8 @@ interface = gr.Interface(
     Model Repository: {MODEL_REPO}
     Current User: {CURRENT_USER}
     Last Updated: {CURRENT_UTC} UTC
+    
+    Enter your question and click submit to get a response.
     """,
     theme=gr.themes.Soft(),
     examples=[