Spaces:
Running
Running
| """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 | |