Spaces:
Sleeping
Sleeping
gary-boon
Claude
commited on
Commit
·
f5ba954
1
Parent(s):
9e9dc34
Fix research attention endpoint model compatibility
Browse filesMake 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>
- 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 |
-
|
| 1332 |
-
|
| 1333 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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"]
|