Fix: Refine layer hook output format handling

- Simplified logic to match exact output structure
- Ensure compatibility with layer_norm expectations
- Handle all tuple/tensor cases properly

Testing different approach to prevent layer_norm type errors.

🤖 Generated with [Claude Code](https://claude.ai/code)

Co-Authored-By: Claude <noreply@anthropic.com>

backend/model_service.py

@@ -298,29 +298,27 @@ class ModelManager:
             
             def create_layer_hook():
                 def hook(module, input, output):
-                    # Skip layer by passing through input unchanged
-                    # The input to a transformer layer is (hidden_states, optional_attention_mask, ...)
-                    # The output is (hidden_states, optional_attention_weights, ...)
-                    # We want to pass the input hidden states as if the layer did nothing
+                    # Skip layer by making it an identity operation
+                    # The key insight: we must match the EXACT output structure
+                    # but replace hidden states with input hidden states
                     
-                    # Get the input hidden states
-                    if isinstance(input, tuple) and len(input) > 0:
-                        input_hidden_states = input[0]
-                    else:
-                        input_hidden_states = input
+                    # For CodeGen blocks, the input/output structure is:
+                    # input: (hidden_states,) or just hidden_states
+                    # output: (hidden_states,) or (hidden_states, presents) etc.
                     
-                    # Return in the same format as the output
-                    if isinstance(output, tuple):
-                        # Check if there are additional elements to preserve
-                        if len(output) > 1:
-                            # Keep any additional outputs (like attention weights)
-                            return (input_hidden_states,) + output[1:]
-                        else:
-                            # Output is a single-element tuple, return the same
-                            return (input_hidden_states,)
-                    else:
-                        # Output is a plain tensor, return input as plain tensor
+                    # Get input hidden states
+                    input_hidden_states = input[0] if isinstance(input, tuple) else input
+                    
+                    # Match output structure exactly
+                    if not isinstance(output, tuple):
+                        # If output is a plain tensor, return input as plain tensor
                         return input_hidden_states
+                    elif len(output) == 1:
+                        # Single element tuple - preserve as single element tuple
+                        return (input_hidden_states,)
+                    else:
+                        # Multiple elements - keep all but replace hidden states
+                        return (input_hidden_states,) + output[1:]
                 return hook
             
             # Apply hooks and log what's being disabled