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:54:34"
+CURRENT_UTC = "2025-05-06 16:00:41"
 
 def setup_logging():
     logging.basicConfig(
@@ -28,16 +28,15 @@ class PositionalEncoding(torch.nn.Module):
         super().__init__()
         self.dropout = torch.nn.Dropout(p=dropout)
         
-        pe = torch.zeros(max_len, d_model)
+        pe = torch.zeros(max_len, 1, d_model)  # Changed dimension order to match saved 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)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
         self.register_buffer('pe', pe)
         
     def forward(self, x):
-        x = x + self.pe[:, :x.size(1)]
+        x = x + self.pe[:x.size(0)]
         return self.dropout(x)
 
 class ImprovedTransformer(torch.nn.Module):
@@ -85,9 +84,15 @@ class ImprovedTransformer(torch.nn.Module):
         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))
         
+        src = src.transpose(0, 1)  # Change to time-first
+        tgt = tgt.transpose(0, 1)  # Change to time-first
+        
         src = self.pos_encoder(src)
         tgt = self.pos_encoder(tgt)
         
+        src = src.transpose(0, 1)  # Back to batch-first
+        tgt = tgt.transpose(0, 1)  # Back to batch-first
+        
         # Transform
         output = self.transformer(
             src,
@@ -128,7 +133,7 @@ class ModelInference:
                 token=token
             )
             
-            # Load model data first to get configuration
+            # Load model data first
             self.logger.info("Loading model data...")
             model_data = torch.load(
                 model_path,
@@ -140,20 +145,21 @@ class ModelInference:
             with open(tokenizer_path, 'r', encoding='utf-8') as f:
                 tokenizer_data = json.load(f)
             
-            # Initialize tokenizer with the same vocabulary size as the saved model
+            # Get exact vocabulary size from the saved model
             self.vocab = tokenizer_data['vocab']
-            vocab_size = len(self.vocab)  # Use exact vocab size from saved model
+            vocab_size = 1747  # Exact size from the saved model
             
+            # Initialize special tokens to match the saved model
             self.special_tokens = {
-                "<user>": vocab_size,
-                "<assistant>": vocab_size + 1,
-                "<sep>": vocab_size + 2,
-                "<eos>": vocab_size + 3
+                "<user>": vocab_size - 4,
+                "<assistant>": vocab_size - 3,
+                "<sep>": vocab_size - 2,
+                "<eos>": vocab_size - 1
             }
             
-            # Initialize model with exact vocab size
+            # Initialize model with exact vocab size from saved model
             self.model = ImprovedTransformer(
-                vocab_size=vocab_size + 4,  # Add exactly 4 special tokens
+                vocab_size=vocab_size,  # Use exact size
                 d_model=512,
                 nhead=8,
                 num_encoder_layers=3,
@@ -161,19 +167,8 @@ 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(fixed_state_dict, strict=True)
+            self.model.load_state_dict(model_data['model_state_dict'])
             self.model.eval()
             
             self.logger.info("Model loaded successfully")