adds POST /tokenize endpoint

Author: anthonyduong9

apps/inference/neuronpedia_inference/endpoints/tokenize.py

@@ -0,0 +1,59 @@
+import logging
+
+from fastapi import APIRouter
+from fastapi.responses import JSONResponse
+from neuronpedia_inference_client.models.tokenize_post200_response import (
+    TokenizePost200Response,
+)
+from neuronpedia_inference_client.models.tokenize_post_request import (
+    TokenizePostRequest,
+)
+
+from neuronpedia_inference.config import Config
+from neuronpedia_inference.shared import Model, with_request_lock
+
+logger = logging.getLogger(__name__)
+
+router = APIRouter()
+
+
+@router.post("/tokenize")
+@with_request_lock()
+async def tokenize(
+    request: TokenizePostRequest,
+):
+    model = Model.get_instance()
+    config = Config.get_instance()
+
+    prepend_bos = (
+        request.prepend_bos
+        if request.prepend_bos is not None
+        else model.cfg.default_prepend_bos
+    )
+
+    tokens = model.to_tokens(
+        request.text,
+        prepend_bos=prepend_bos,
+        truncate=False,
+    )[0]
+
+    if len(tokens) > config.TOKEN_LIMIT:
+        logger.error(
+            "Text too long: %s tokens, max is %s",
+            len(tokens),
+            config.TOKEN_LIMIT,
+        )
+        return JSONResponse(
+            content={
+                "error": f"Text too long: {len(tokens)} tokens, max is {config.TOKEN_LIMIT}"
+            },
+            status_code=400,
+        )
+
+    str_tokens = model.to_str_tokens(request.text, prepend_bos=prepend_bos)
+
+    return TokenizePost200Response(
+        tokens=tokens.tolist(),
+        token_strings=str_tokens,  # type: ignore
+        prepend_bos=prepend_bos,
+    )

apps/inference/neuronpedia_inference/server.py

@@ -33,6 +33,9 @@
 from neuronpedia_inference.endpoints.steer.completion_chat import (
     router as steer_completion_chat_router,
 )
+from neuronpedia_inference.endpoints.tokenize import (
+    router as tokenize_router,
+)
 from neuronpedia_inference.endpoints.util.sae_topk_by_decoder_cossim import (
     router as sae_topk_by_decoder_cossim_router,
 )
@@ -88,6 +91,7 @@ async def startup_event():
 v1_router.include_router(activation_topk_by_token_router)
 v1_router.include_router(sae_topk_by_decoder_cossim_router)
 v1_router.include_router(sae_vector_router)
+v1_router.include_router(tokenize_router)
 
 app.include_router(v1_router)
 

apps/inference/poetry.lock

@@ -28,8 +28,8 @@ sae-lens = "^5.8.1"
 fastapi = "^0.115.6"
 uvicorn = "^0.34.0"
 sentry-sdk = {extras = ["fastapi"], version = "^2.20.0"}
-neuronpedia-inference-client = {path = "../../packages/python/neuronpedia-inference-client"}
 nnsight = "^0.4.3"
+neuronpedia-inference-client = {path = "../../packages/python/neuronpedia-inference-client"}
 
 [tool.poetry.group.dev.dependencies]
 # line-profiler = "^4.1.3"

apps/inference/pyproject.toml

@@ -28,4 +28,4 @@ jobs:
           pip install -r test-requirements.txt
       - name: Test with pytest
         run: |
-          pytest --cov={{packageName}}
+          pytest --cov=neuronpedia_inference_client

packages/python/neuronpedia-inference-client/.github/workflows/python.yml

@@ -27,6 +27,8 @@ docs/SteerCompletionChatPost200Response.md
 docs/SteerCompletionChatPostRequest.md
 docs/SteerCompletionPost200Response.md
 docs/SteerCompletionRequest.md
+docs/TokenizePost200Response.md
+docs/TokenizePostRequest.md
 docs/UtilSaeTopkByDecoderCossimPost200Response.md
 docs/UtilSaeTopkByDecoderCossimPost200ResponseTopkDecoderCossimFeaturesInner.md
 docs/UtilSaeTopkByDecoderCossimPostRequest.md
@@ -63,6 +65,8 @@ neuronpedia_inference_client/models/steer_completion_chat_post200_response.py
 neuronpedia_inference_client/models/steer_completion_chat_post_request.py
 neuronpedia_inference_client/models/steer_completion_post200_response.py
 neuronpedia_inference_client/models/steer_completion_request.py
+neuronpedia_inference_client/models/tokenize_post200_response.py
+neuronpedia_inference_client/models/tokenize_post_request.py
 neuronpedia_inference_client/models/util_sae_topk_by_decoder_cossim_post200_response.py
 neuronpedia_inference_client/models/util_sae_topk_by_decoder_cossim_post200_response_topk_decoder_cossim_features_inner.py
 neuronpedia_inference_client/models/util_sae_topk_by_decoder_cossim_post_request.py
@@ -99,6 +103,8 @@ test/test_steer_completion_chat_post200_response.py
 test/test_steer_completion_chat_post_request.py
 test/test_steer_completion_post200_response.py
 test/test_steer_completion_request.py
+test/test_tokenize_post200_response.py
+test/test_tokenize_post_request.py
 test/test_util_sae_topk_by_decoder_cossim_post200_response.py
 test/test_util_sae_topk_by_decoder_cossim_post200_response_topk_decoder_cossim_features_inner.py
 test/test_util_sae_topk_by_decoder_cossim_post_request.py

packages/python/neuronpedia-inference-client/.openapi-generator/FILES

@@ -1 +1 @@
-7.11.0
+7.12.0

packages/python/neuronpedia-inference-client/.openapi-generator/VERSION

@@ -4,8 +4,8 @@ No description provided (generated by Openapi Generator https://github.com/opena
 This Python package is automatically generated by the [OpenAPI Generator](https://openapi-generator.tech) project:
 
 - API version: 1.1.0
-- Package version: 1.1.1
-- Generator version: 7.11.0
+- Package version: 1.2.0
+- Generator version: 7.12.0
 - Build package: org.openapitools.codegen.languages.PythonClientCodegen
 
 ## Requirements.
@@ -100,6 +100,7 @@ Class | Method | HTTP request | Description
 *DefaultApi* | [**activation_topk_by_token_post**](docs/DefaultApi.md#activation_topk_by_token_post) | **POST** /activation/topk-by-token | For a given prompt, get the top activating features at each token position for a single SAE.
 *DefaultApi* | [**steer_completion_chat_post**](docs/DefaultApi.md#steer_completion_chat_post) | **POST** /steer/completion-chat | For a given prompt, complete it by steering with the given feature or vector
 *DefaultApi* | [**steer_completion_post**](docs/DefaultApi.md#steer_completion_post) | **POST** /steer/completion | For a given prompt, complete it by steering with the given feature or vector
+*DefaultApi* | [**tokenize_post**](docs/DefaultApi.md#tokenize_post) | **POST** /tokenize | Tokenize input text for a given model
 *DefaultApi* | [**util_sae_topk_by_decoder_cossim_post**](docs/DefaultApi.md#util_sae_topk_by_decoder_cossim_post) | **POST** /util/sae-topk-by-decoder-cossim | Given a specific vector or SAE feature, return the top features by cosine similarity in the same SAE
 *DefaultApi* | [**util_sae_vector_post**](docs/DefaultApi.md#util_sae_vector_post) | **POST** /util/sae-vector | Get the raw vector for an SAE feature
 
@@ -128,6 +129,8 @@ Class | Method | HTTP request | Description
  - [SteerCompletionChatPostRequest](docs/SteerCompletionChatPostRequest.md)
  - [SteerCompletionPost200Response](docs/SteerCompletionPost200Response.md)
  - [SteerCompletionRequest](docs/SteerCompletionRequest.md)
+ - [TokenizePost200Response](docs/TokenizePost200Response.md)
+ - [TokenizePostRequest](docs/TokenizePostRequest.md)
  - [UtilSaeTopkByDecoderCossimPost200Response](docs/UtilSaeTopkByDecoderCossimPost200Response.md)
  - [UtilSaeTopkByDecoderCossimPost200ResponseTopkDecoderCossimFeaturesInner](docs/UtilSaeTopkByDecoderCossimPost200ResponseTopkDecoderCossimFeaturesInner.md)
  - [UtilSaeTopkByDecoderCossimPostRequest](docs/UtilSaeTopkByDecoderCossimPostRequest.md)

packages/python/neuronpedia-inference-client/README.md

@@ -9,6 +9,7 @@ Method | HTTP request | Description
 [**activation_topk_by_token_post**](DefaultApi.md#activation_topk_by_token_post) | **POST** /activation/topk-by-token | For a given prompt, get the top activating features at each token position for a single SAE.
 [**steer_completion_chat_post**](DefaultApi.md#steer_completion_chat_post) | **POST** /steer/completion-chat | For a given prompt, complete it by steering with the given feature or vector
 [**steer_completion_post**](DefaultApi.md#steer_completion_post) | **POST** /steer/completion | For a given prompt, complete it by steering with the given feature or vector
+[**tokenize_post**](DefaultApi.md#tokenize_post) | **POST** /tokenize | Tokenize input text for a given model
 [**util_sae_topk_by_decoder_cossim_post**](DefaultApi.md#util_sae_topk_by_decoder_cossim_post) | **POST** /util/sae-topk-by-decoder-cossim | Given a specific vector or SAE feature, return the top features by cosine similarity in the same SAE
 [**util_sae_vector_post**](DefaultApi.md#util_sae_vector_post) | **POST** /util/sae-vector | Get the raw vector for an SAE feature
 
@@ -408,6 +409,85 @@ Name | Type | Description  | Notes
 
 [[Back to top]](#) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to Model list]](../README.md#documentation-for-models) [[Back to README]](../README.md)
 
+# **tokenize_post**
+> TokenizePost200Response tokenize_post(tokenize_post_request)
+
+Tokenize input text for a given model
+
+### Example
+
+* Api Key Authentication (SimpleSecretAuth):
+
+```python
+import neuronpedia_inference_client
+from neuronpedia_inference_client.models.tokenize_post200_response import TokenizePost200Response
+from neuronpedia_inference_client.models.tokenize_post_request import TokenizePostRequest
+from neuronpedia_inference_client.rest import ApiException
+from pprint import pprint
+
+# Defining the host is optional and defaults to /v1
+# See configuration.py for a list of all supported configuration parameters.
+configuration = neuronpedia_inference_client.Configuration(
+    host = "/v1"
+)
+
+# The client must configure the authentication and authorization parameters
+# in accordance with the API server security policy.
+# Examples for each auth method are provided below, use the example that
+# satisfies your auth use case.
+
+# Configure API key authorization: SimpleSecretAuth
+configuration.api_key['SimpleSecretAuth'] = os.environ["API_KEY"]
+
+# Uncomment below to setup prefix (e.g. Bearer) for API key, if needed
+# configuration.api_key_prefix['SimpleSecretAuth'] = 'Bearer'
+
+# Enter a context with an instance of the API client
+with neuronpedia_inference_client.ApiClient(configuration) as api_client:
+    # Create an instance of the API class
+    api_instance = neuronpedia_inference_client.DefaultApi(api_client)
+    tokenize_post_request = neuronpedia_inference_client.TokenizePostRequest() # TokenizePostRequest | 
+
+    try:
+        # Tokenize input text for a given model
+        api_response = api_instance.tokenize_post(tokenize_post_request)
+        print("The response of DefaultApi->tokenize_post:\n")
+        pprint(api_response)
+    except Exception as e:
+        print("Exception when calling DefaultApi->tokenize_post: %s\n" % e)
+```
+
+
+
+### Parameters
+
+
+Name | Type | Description  | Notes
+------------- | ------------- | ------------- | -------------
+ **tokenize_post_request** | [**TokenizePostRequest**](TokenizePostRequest.md)|  | 
+
+### Return type
+
+[**TokenizePost200Response**](TokenizePost200Response.md)
+
+### Authorization
+
+[SimpleSecretAuth](../README.md#SimpleSecretAuth)
+
+### HTTP request headers
+
+ - **Content-Type**: application/json
+ - **Accept**: application/json
+
+### HTTP response details
+
+| Status code | Description | Response headers |
+|-------------|-------------|------------------|
+**200** | Successful tokenization |  -  |
+**401** | X-SECRET-KEY header is missing or invalid |  * WWW_Authenticate -  <br>  |
+
+[[Back to top]](#) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to Model list]](../README.md#documentation-for-models) [[Back to README]](../README.md)
+
 # **util_sae_topk_by_decoder_cossim_post**
 > UtilSaeTopkByDecoderCossimPost200Response util_sae_topk_by_decoder_cossim_post(util_sae_topk_by_decoder_cossim_post_request)
 

packages/python/neuronpedia-inference-client/docs/DefaultApi.md

@@ -0,0 +1,31 @@
+# TokenizePost200Response
+
+
+## Properties
+
+Name | Type | Description | Notes
+------------ | ------------- | ------------- | -------------
+**tokens** | **List[int]** | Token IDs for the input text | 
+**token_strings** | **List[str]** | String representation of each token | 
+**prepend_bos** | **bool** | Whether beginning-of-sequence token was prepended | 
+
+## Example
+
+```python
+from neuronpedia_inference_client.models.tokenize_post200_response import TokenizePost200Response
+
+# TODO update the JSON string below
+json = "{}"
+# create an instance of TokenizePost200Response from a JSON string
+tokenize_post200_response_instance = TokenizePost200Response.from_json(json)
+# print the JSON string representation of the object
+print(TokenizePost200Response.to_json())
+
+# convert the object into a dict
+tokenize_post200_response_dict = tokenize_post200_response_instance.to_dict()
+# create an instance of TokenizePost200Response from a dict
+tokenize_post200_response_from_dict = TokenizePost200Response.from_dict(tokenize_post200_response_dict)
+```
+[[Back to Model list]](../README.md#documentation-for-models) [[Back to API list]](../README.md#documentation-for-api-endpoints) [[Back to README]](../README.md)
+
+