Update modeling_modified.py

modeling_modified.py

@@ -12,13 +12,30 @@ class ModifiedBartForConditionalGenerationWithHook(BartForConditionalGeneration)
     def forward(self, *args, **kwargs):
         # Log input information
         input_ids = kwargs.get('input_ids', args[0] if args else None)
-        print(f"Forward pass initiated with input shape: {input_ids.shape if input_ids is not None else 'None'}")
-        
+        logger.info(f"Forward pass initiated with input shape: {input_ids.shape if input_ids is not None else 'None'}")
+
         # Call the parent's forward method
         outputs = super().forward(*args, **kwargs)
-        
+
+        # Log output information
+        logger.info(f"Forward pass completed. Output shape: {outputs.logits.shape if hasattr(outputs, 'logits') else 'N/A'}")
+
+        return outputs
+
+    def generate(self, *args, **kwargs):
+        if not hasattr(super(), 'generate'):
+            logger.warning("Generate method is not available in the parent class.")
+            return None
+
+        # Log input information
+        input_ids = kwargs.get('input_ids', args[0] if args else None)
+        logger.info(f"Generate method called with input shape: {input_ids.shape if input_ids is not None else 'None'}")
+
+        # Call the parent's generate method
+        outputs = super().generate(*args, **kwargs)
+
         # Log output information
-        print(f"Forward pass completed. Output shape: {outputs.logits.shape if hasattr(outputs, 'logits') else 'N/A'}")
-        
+        logger.info(f"Generate method completed. Output shape: {outputs.shape}")
+
         return outputs