TokenTrace / backend /api /activation_explain.py
cccmmd
feat: add Tiny-NLA activation explanation with trained model weights
9d0d4e9
Raw
History Blame Contribute Delete
4.13 kB
"""Activation Explainer API (Tiny-NLA)"""
import gc
import time
import torch
from backend.platform.oom import exit_if_oom
from backend.api.analyze import LOCK_WAIT_TIMEOUT
from backend.platform.access_log import get_client_ip, log_prediction_attribute_request
from backend.platform.source_page import ALLOWED_SOURCE_PAGES, normalize_source_page
from backend.core.tiny_nla import TinyNLAEngine, tiny_nla_lock, TINY_NLA_LOCK_TIMEOUT
def activation_explain(activation_explain_request):
model = activation_explain_request.get("model")
source_page = activation_explain_request.get("source_page")
text = activation_explain_request.get("text")
token_index = activation_explain_request.get("token_index")
vector = activation_explain_request.get("vector")
if model is None:
return {"success": False, "message": "Missing required field: model"}, 400
if model not in ("base", "instruct"):
return {"success": False, "message": 'model must be "base" or "instruct"'}, 400
if source_page is None or source_page == "":
return {"success": False, "message": "Missing required field: source_page"}, 400
normalized_source_page = normalize_source_page(source_page)
if normalized_source_page is None:
allowed = ", ".join(sorted(ALLOWED_SOURCE_PAGES))
return {"success": False, "message": f"source_page must be one of: {allowed}"}, 400
source_page = normalized_source_page
has_text = isinstance(text, str) and text.strip() != ""
has_vector = isinstance(vector, list) and len(vector) > 0
if not has_text and not has_vector:
return {"success": False, "message": "Missing required field: text or vector must be provided"}, 400
if token_index is not None and not isinstance(token_index, int):
return {"success": False, "message": "token_index must be an integer"}, 400
if token_index is not None and token_index < 0:
return {"success": False, "message": "token_index must be >= 0"}, 400
if has_text and token_index is None:
return {"success": False, "message": "token_index is required when text is provided"}, 400
client_ip = get_client_ip()
start_time = time.perf_counter()
request_id = log_prediction_attribute_request(
context=text if has_text else str(vector)[:200],
target_prediction=None,
target_token_id=token_index,
model=model,
source_page=source_page,
flow_id=None,
flow_step=None,
client_ip=client_ip,
)
lock_acquired = tiny_nla_lock.acquire(timeout=TINY_NLA_LOCK_TIMEOUT)
if not lock_acquired:
return {"success": False, "message": f"Tiny-NLA queue wait exceeded {TINY_NLA_LOCK_TIMEOUT} seconds; please try again later."}, 503
try:
engine = TinyNLAEngine()
if has_vector:
activation = torch.tensor(vector, dtype=torch.float32)
if activation.shape[0] != 1024:
return {"success": False, "message": f"vector dimension must be 1024, got {activation.shape[0]}"}, 400
else:
activation = engine.extract_activation(text, token_index)
explanation = engine.explain(activation)
roundtrip_cosine = engine.reconstruct_cosine(activation, explanation)
result = {
"concept": "",
"explanation": explanation,
"roundtrip_cosine": round(roundtrip_cosine, 4),
"vector_dim": 1024,
"note": "",
}
except ValueError as e:
return {"success": False, "message": str(e)}, 400
except Exception as e:
import traceback
traceback.print_exc()
exit_if_oom(e, defer_seconds=1)
return {"success": False, "message": str(e)}, 500
finally:
tiny_nla_lock.release()
gc.collect()
elapsed = time.perf_counter() - start_time
print(
f"\t📤 API activation_explain response: req_id={request_id}, "
f"concept={result.get('concept')!r}, roundtrip={result.get('roundtrip_cosine')}, "
f"response_time={elapsed:.4f}s"
)
return {"success": True, **result}, 200