gary-boon Claude commited on
Commit
f5ba954
·
1 Parent(s): 9e9dc34

Fix research attention endpoint model compatibility

Browse files

Make QKV hook registration more robust to support different model
architectures (CodeGen vs GPT-2 style attention modules)

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

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

Files changed (1) hide show
  1. backend/model_service.py +14 -4
backend/model_service.py CHANGED
@@ -1327,10 +1327,20 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
1327
  }
1328
  return hook
1329
 
1330
- # Register hooks on all qkv_proj modules
1331
- for layer_idx, layer in enumerate(manager.model.transformer.h):
1332
- hook = layer.attn.qkv_proj.register_forward_hook(make_qkv_hook(layer_idx))
1333
- hooks.append(hook)
 
 
 
 
 
 
 
 
 
 
1334
 
1335
  with torch.no_grad():
1336
  current_ids = inputs["input_ids"]
 
1327
  }
1328
  return hook
1329
 
1330
+ # Register hooks on all qkv_proj modules (if available)
1331
+ # This is model-specific - CodeGen uses different architecture
1332
+ try:
1333
+ if hasattr(manager.model, 'transformer') and hasattr(manager.model.transformer, 'h'):
1334
+ for layer_idx, layer in enumerate(manager.model.transformer.h):
1335
+ if hasattr(layer, 'attn') and hasattr(layer.attn, 'qkv_proj'):
1336
+ hook = layer.attn.qkv_proj.register_forward_hook(make_qkv_hook(layer_idx))
1337
+ hooks.append(hook)
1338
+ elif hasattr(layer, 'attn') and hasattr(layer.attn, 'c_attn'):
1339
+ # GPT-2 style attention
1340
+ hook = layer.attn.c_attn.register_forward_hook(make_qkv_hook(layer_idx))
1341
+ hooks.append(hook)
1342
+ except Exception as hook_error:
1343
+ logger.warning(f"Could not register QKV hooks: {hook_error}")
1344
 
1345
  with torch.no_grad():
1346
  current_ids = inputs["input_ids"]