crfm-helm 0.4.0__py3-none-any.whl → 0.5.1__py3-none-any.whl

helm/proxy/server.py

@@ -1,5 +1,3 @@
-# mypy: check_untyped_defs = False
-
 """
 Starts a REST server for the frontend to interact with.
 Look at `index.js` to see how the functionality is invoked.
@@ -20,12 +18,17 @@ from helm.benchmark.config_registry import (
     register_configs_from_directory,
     register_builtin_configs_from_helm_package,
 )
+from helm.benchmark.model_deployment_registry import get_default_model_deployment_for_model
 from helm.common.authentication import Authentication
+from helm.common.cache_backend_config import CacheBackendConfig, MongoCacheBackendConfig, SqliteCacheBackendConfig
+from helm.common.general import ensure_directory_exists
 from helm.common.hierarchical_logger import hlog
 from helm.common.optional_dependencies import handle_module_not_found_error
 from helm.common.request import Request
 from helm.common.perspective_api_request import PerspectiveAPIRequest
+from helm.common.moderations_api_request import ModerationAPIRequest
 from helm.common.tokenization_request import TokenizationRequest, DecodeRequest
+from helm.proxy.services.service import CACHE_DIR
 from .accounts import Account
 from .services.server_service import ServerService
 from .query import Query
@@ -39,6 +42,7 @@ except ModuleNotFoundError as e:
 bottle.BaseRequest.MEMFILE_MAX = 1024 * 1024
 
 app = bottle.default_app()
+service: ServerService
 
 
 def safe_call(func, to_json=True):
@@ -87,9 +91,16 @@ def handle_static_filename(filename):
     return resp
 
 
+@app.get("/output/<filename:path>")
+def handle_output_filename(filename):
+    resp = bottle.static_file(filename, root=app.config["crfm.proxy.outputpath"])
+    return resp
+
+
 @app.get("/api/general_info")
 def handle_get_general_info():
     def perform(args):
+        global service
         return dataclasses.asdict(service.get_general_info())
 
     return safe_call(perform)
@@ -98,6 +109,7 @@ def handle_get_general_info():
 @app.get("/api/window_service_info")
 def handle_get_window_service_info():
     def perform(args):
+        global service
         return dataclasses.asdict(service.get_window_service_info(args["model_name"]))
 
     return safe_call(perform)
@@ -106,6 +118,7 @@ def handle_get_window_service_info():
 @app.post("/api/account")
 def handle_create_account():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         return dataclasses.asdict(service.create_account(auth))
 
@@ -115,6 +128,7 @@ def handle_create_account():
 @app.delete("/api/account")
 def handle_delete_account():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         api_key = args["api_key"]
         return dataclasses.asdict(service.delete_account(auth, api_key))
@@ -125,6 +139,7 @@ def handle_delete_account():
 @app.get("/api/account")
 def handle_get_account():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         if "all" in args and args["all"].lower() == "true":
             return [dataclasses.asdict(account) for account in service.get_accounts(auth)]
@@ -137,6 +152,7 @@ def handle_get_account():
 @app.put("/api/account")
 def handle_update_account():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         account = from_dict(Account, json.loads(args["account"]))
         return dataclasses.asdict(service.update_account(auth, account))
@@ -147,6 +163,7 @@ def handle_update_account():
 @app.put("/api/account/api_key")
 def handle_update_api_key():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         account = from_dict(Account, json.loads(args["account"]))
         return dataclasses.asdict(service.rotate_api_key(auth, account))
@@ -157,6 +174,7 @@ def handle_update_api_key():
 @app.get("/api/query")
 def handle_query():
     def perform(args):
+        global service
         query = Query(**args)
         return dataclasses.asdict(service.expand_query(query))
 
@@ -166,9 +184,28 @@ def handle_query():
 @app.get("/api/request")
 def handle_request():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         request = Request(**json.loads(args["request"]))
-        return dataclasses.asdict(service.make_request(auth, request))
+        # Hack to maintain reverse compatibility with clients with version <= 0.3.0.
+        # Clients with version <= 0.3.0 do not set model_deployment, but this is now
+        # required by Request.
+        if not request.model_deployment:
+            model_deployment = get_default_model_deployment_for_model(request.model)
+            if model_deployment is None:
+                raise ValueError(f"Unknown model '{request.model}'")
+            request = dataclasses.replace(request, model_deployment=model_deployment)
+
+        raw_response = dataclasses.asdict(service.make_request(auth, request))
+
+        # Hack to maintain reverse compatibility with clients with version <= 1.0.0.
+        # Clients with version <= 1.0.0 expect each token to contain a `top_logprobs`
+        # field of type dict.
+        for completion in raw_response["completions"]:
+            for token in completion["tokens"]:
+                token["top_logprobs"] = {}
+
+        return raw_response
 
     return safe_call(perform)
 
@@ -176,6 +213,7 @@ def handle_request():
 @app.get("/api/tokenize")
 def handle_tokenization():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         request = TokenizationRequest(**json.loads(args["request"]))
         return dataclasses.asdict(service.tokenize(auth, request))
@@ -186,6 +224,7 @@ def handle_tokenization():
 @app.get("/api/decode")
 def handle_decode():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         request = DecodeRequest(**json.loads(args["request"]))
         return dataclasses.asdict(service.decode(auth, request))
@@ -196,6 +235,7 @@ def handle_decode():
 @app.get("/api/toxicity")
 def handle_toxicity_request():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         request = PerspectiveAPIRequest(**json.loads(args["request"]))
         return dataclasses.asdict(service.get_toxicity_scores(auth, request))
@@ -203,9 +243,21 @@ def handle_toxicity_request():
     return safe_call(perform)
 
 
+@app.get("/api/moderation")
+def handle_moderation_request():
+    def perform(args):
+        global service
+        auth = Authentication(**json.loads(args["auth"]))
+        request = ModerationAPIRequest(**json.loads(args["request"]))
+        return dataclasses.asdict(service.get_moderation_results(auth, request))
+
+    return safe_call(perform)
+
+
 @app.get("/api/shutdown")
 def handle_shutdown():
     def perform(args):
+        global service
         auth = Authentication(**json.loads(args["auth"]))
         service.shutdown(auth)
 
@@ -218,6 +270,7 @@ def main():
     parser.add_argument("-p", "--port", type=int, help="What port to listen on", default=1959)
     parser.add_argument("--ssl-key-file", type=str, help="Path to SSL key file")
     parser.add_argument("--ssl-cert-file", type=str, help="Path to SSL cert file")
+    parser.add_argument("--ssl-ca-certs", type=str, help="Path to SSL CA certs")
     parser.add_argument("-b", "--base-path", help="What directory has credentials, etc.", default="prod_env")
     parser.add_argument("-w", "--workers", type=int, help="Number of worker processes to handle requests", default=8)
     parser.add_argument("-t", "--timeout", type=int, help="Request timeout in seconds", default=5 * 60)
@@ -232,17 +285,29 @@ def main():
     register_builtin_configs_from_helm_package()
     register_configs_from_directory(args.base_path)
 
-    service = ServerService(base_path=args.base_path, mongo_uri=args.mongo_uri)
+    cache_backend_config: CacheBackendConfig
+    if args.mongo_uri:
+        cache_backend_config = MongoCacheBackendConfig(args.mongo_uri)
+    else:
+        sqlite_cache_path = os.path.join(args.base_path, CACHE_DIR)
+        ensure_directory_exists(sqlite_cache_path)
+        cache_backend_config = SqliteCacheBackendConfig(sqlite_cache_path)
+
+    service = ServerService(base_path=args.base_path, cache_backend_config=cache_backend_config)
 
     gunicorn_args = {
         "workers": args.workers,
         "timeout": args.timeout,
         "limit_request_line": 0,  # Controls the maximum size of HTTP request line in bytes. 0 = unlimited.
     }
-    if args.ssl_key_file and args.ssl_cert_file:
+    if args.ssl_key_file:
         gunicorn_args["keyfile"] = args.ssl_key_file
+    if args.ssl_cert_file:
         gunicorn_args["certfile"] = args.ssl_cert_file
+    if args.ssl_ca_certs:
+        gunicorn_args["ca_certs"] = args.ssl_ca_certs
 
     # Clear arguments before running gunicorn as it also uses argparse
     sys.argv = [sys.argv[0]]
+    app.config["crfm.proxy.outputpath"] = os.path.join(os.path.realpath(args.base_path), "cache", "output")
     app.run(host="0.0.0.0", port=args.port, server="gunicorn", **gunicorn_args)

helm/proxy/services/remote_service.py

@@ -5,9 +5,15 @@ import urllib.parse
 from dataclasses import asdict
 from typing import Any, List, Optional
 
+from helm.common.cache import CacheConfig
+from helm.common.cache_backend_config import BlackHoleCacheBackendConfig
 from helm.common.authentication import Authentication
+from helm.common.moderations_api_request import ModerationAPIRequest, ModerationAPIRequestResult
 from helm.common.critique_request import CritiqueRequest, CritiqueRequestResult
+from helm.common.nudity_check_request import NudityCheckRequest, NudityCheckResult
+from helm.common.file_upload_request import FileUploadRequest, FileUploadResult
 from helm.common.perspective_api_request import PerspectiveAPIRequest, PerspectiveAPIRequestResult
+from helm.common.clip_score_request import CLIPScoreRequest, CLIPScoreResult
 from helm.common.tokenization_request import (
     WindowServiceInfo,
     TokenizationRequest,
@@ -27,6 +33,8 @@ class RemoteServiceError(Exception):
 
 
 class RemoteService(Service):
+    NOT_SUPPORTED_ERROR: str = "Not supported through the remote service."
+
     def __init__(self, base_url):
         self.base_url: str = base_url
 
@@ -84,6 +92,15 @@ class RemoteService(Service):
         RemoteService._check_response(response, request_json)
         return from_dict(DecodeRequestResult, response)
 
+    def upload(self, auth: Authentication, request: FileUploadRequest) -> FileUploadResult:
+        raise NotImplementedError(self.NOT_SUPPORTED_ERROR)
+
+    def check_nudity(self, auth: Authentication, request: NudityCheckRequest) -> NudityCheckResult:
+        raise NotImplementedError(self.NOT_SUPPORTED_ERROR)
+
+    def compute_clip_score(self, auth: Authentication, request: CLIPScoreRequest) -> CLIPScoreResult:
+        raise NotImplementedError(self.NOT_SUPPORTED_ERROR)
+
     def get_toxicity_scores(self, auth: Authentication, request: PerspectiveAPIRequest) -> PerspectiveAPIRequestResult:
         request_json: str = json.dumps(asdict(request))
         params = {
@@ -94,6 +111,16 @@ class RemoteService(Service):
         RemoteService._check_response(response, request_json)
         return from_dict(PerspectiveAPIRequestResult, response)
 
+    def get_moderation_results(self, auth: Authentication, request: ModerationAPIRequest) -> ModerationAPIRequestResult:
+        request_json: str = json.dumps(asdict(request))
+        params = {
+            "auth": json.dumps(asdict(auth)),
+            "request": request_json,
+        }
+        response = requests.get(f"{self.base_url}/api/moderation?{urllib.parse.urlencode(params)}").json()
+        RemoteService._check_response(response, request_json)
+        return from_dict(ModerationAPIRequestResult, response)
+
     def make_critique_request(self, auth: Authentication, request: CritiqueRequest) -> CritiqueRequestResult:
         raise NotImplementedError("make_critique_request is not supported by RemoteServer")
 
@@ -153,6 +180,10 @@ class RemoteService(Service):
             # A ConnectionError is expected when shutting down the server.
             pass
 
+    def get_cache_config(self, shard_name: str) -> CacheConfig:
+        """Returns a CacheConfig"""
+        return BlackHoleCacheBackendConfig().get_cache_config(shard_name)
+
 
 def add_service_args(parser: argparse.ArgumentParser):
     """Add command-line arguments to enable command-line utilities to specify how to connect to a remote server."""

helm/proxy/services/server_service.py

@@ -2,10 +2,15 @@ import dataclasses
 import os
 import signal
 from typing import List, Optional
-from helm.common.cache_utils import build_cache_config
 
+from helm.common.cache import CacheConfig
+from helm.common.cache_backend_config import CacheBackendConfig, BlackHoleCacheBackendConfig
 from helm.common.critique_request import CritiqueRequest, CritiqueRequestResult
 from helm.common.authentication import Authentication
+from helm.common.moderations_api_request import ModerationAPIRequest, ModerationAPIRequestResult
+from helm.common.clip_score_request import CLIPScoreRequest, CLIPScoreResult
+from helm.common.nudity_check_request import NudityCheckRequest, NudityCheckResult
+from helm.common.file_upload_request import FileUploadRequest, FileUploadResult
 from helm.common.general import ensure_directory_exists, parse_hocon, get_credentials
 from helm.common.perspective_api_request import PerspectiveAPIRequest, PerspectiveAPIRequestResult
 from helm.common.tokenization_request import (
@@ -18,16 +23,20 @@ from helm.common.tokenization_request import (
 from helm.common.request import Request, RequestResult
 from helm.common.hierarchical_logger import hlog
 from helm.proxy.accounts import Accounts, Account
-from helm.proxy.clients.auto_client import AutoClient
-from helm.proxy.clients.toxicity_classifier_client import ToxicityClassifierClient
+from helm.clients.auto_client import AutoClient
+from helm.clients.moderation_api_client import ModerationAPIClient
+from helm.clients.perspective_api_client import PerspectiveAPIClient
+from helm.clients.image_generation.nudity_check_client import NudityCheckClient
+from helm.clients.gcs_client import GCSClient
+from helm.clients.clip_score_client import CLIPScoreClient
+from helm.clients.toxicity_classifier_client import ToxicityClassifierClient
 from helm.proxy.example_queries import example_queries
 from helm.benchmark.model_metadata_registry import ALL_MODELS_METADATA
 from helm.benchmark.model_deployment_registry import get_model_deployment_host_organization
 from helm.proxy.query import Query, QueryResult
 from helm.proxy.retry import retry_request
 from helm.proxy.token_counters.auto_token_counter import AutoTokenCounter
-from helm.proxy.tokenizers.auto_tokenizer import AutoTokenizer
-from helm.proxy.tokenizers.huggingface_tokenizer import HuggingFaceTokenizer
+from helm.tokenizers.auto_tokenizer import AutoTokenizer
 from .service import (
     Service,
     CACHE_DIR,
@@ -44,19 +53,32 @@ class ServerService(Service):
     Main class that supports various functionality for the server.
     """
 
-    def __init__(self, base_path: str = "prod_env", root_mode=False, mongo_uri: str = ""):
+    def __init__(
+        self,
+        base_path: str = "prod_env",
+        root_mode: bool = False,
+        cache_backend_config: CacheBackendConfig = BlackHoleCacheBackendConfig(),
+    ):
+        ensure_directory_exists(base_path)
+        client_file_storage_path = os.path.join(base_path, CACHE_DIR)
+        ensure_directory_exists(client_file_storage_path)
+
         credentials = get_credentials(base_path)
-        cache_path = os.path.join(base_path, CACHE_DIR)
-        ensure_directory_exists(cache_path)
         accounts_path = os.path.join(base_path, ACCOUNTS_FILE)
 
-        self.client = AutoClient(credentials, cache_path, mongo_uri)
-        self.tokenizer = AutoTokenizer(credentials, cache_path, mongo_uri)
-        cache_config = build_cache_config(cache_path, mongo_uri, "huggingface")
-        self.token_counter = AutoTokenCounter(HuggingFaceTokenizer(cache_config=cache_config))
+        self.cache_backend_config = cache_backend_config
+        self.client = AutoClient(credentials, client_file_storage_path, cache_backend_config)
+        self.tokenizer = AutoTokenizer(credentials, cache_backend_config)
+        self.token_counter = AutoTokenCounter(self.tokenizer)
         self.accounts = Accounts(accounts_path, root_mode=root_mode)
-        # Lazily instantiated by get_toxicity_scores()
+
+        # Lazily instantiate the following clients
+        self.moderation_api_client: Optional[ModerationAPIClient] = None
         self.toxicity_classifier_client: Optional[ToxicityClassifierClient] = None
+        self.perspective_api_client: Optional[PerspectiveAPIClient] = None
+        self.nudity_check_client: Optional[NudityCheckClient] = None
+        self.clip_score_client: Optional[CLIPScoreClient] = None
+        self.gcs_client: Optional[GCSClient] = None
 
     def get_general_info(self) -> GeneralInfo:
         # Can't send release_dates in ModelMetadata bacause dates cannot be round-tripped to and from JSON easily.
@@ -91,6 +113,21 @@ class ServerService(Service):
             requests.append(request)
         return QueryResult(requests=requests)
 
+    def _get_model_group_for_model_deployment(self, model_deployment: str) -> str:
+        if model_deployment.startswith("openai/"):
+            if model_deployment.startswith("openai/code-"):
+                return "codex"
+            elif model_deployment.startswith("openai/dall-e-"):
+                return "dall_e"
+            elif model_deployment.startswith("openai/gpt-4-"):
+                return "gpt4"
+            else:
+                return "gpt3"
+        elif model_deployment.startswith("ai21/"):
+            return "jurassic"
+        else:
+            return get_model_deployment_host_organization(model_deployment)
+
     def make_request(self, auth: Authentication, request: Request) -> RequestResult:
         """Actually make a request to an API."""
         # TODO: try to invoke the API even if we're not authenticated, and if
@@ -98,9 +135,9 @@ class ServerService(Service):
         #       https://github.com/stanford-crfm/benchmarking/issues/56
 
         self.accounts.authenticate(auth)
-        host_organization: str = get_model_deployment_host_organization(request.model_deployment)
+        model_group: str = self._get_model_group_for_model_deployment(request.model_deployment)
         # Make sure we can use
-        self.accounts.check_can_use(auth.api_key, host_organization)
+        self.accounts.check_can_use(auth.api_key, model_group)
 
         # Use!
         request_result: RequestResult = self.client.make_request(request)
@@ -109,7 +146,7 @@ class ServerService(Service):
         if not request_result.cached:
             # Count the number of tokens used
             count: int = self.token_counter.count_tokens(request, request_result.completions)
-            self.accounts.use(auth.api_key, host_organization, count)
+            self.accounts.use(auth.api_key, model_group, count)
 
         return request_result
 
@@ -123,6 +160,36 @@ class ServerService(Service):
         self.accounts.authenticate(auth)
         return self.tokenizer.decode(request)
 
+    def upload(self, auth: Authentication, request: FileUploadRequest) -> FileUploadResult:
+        """Uploads a file to external storage."""
+        self.accounts.authenticate(auth)
+
+        if not self.gcs_client:
+            self.gcs_client = self.client.get_gcs_client()
+
+        assert self.gcs_client
+        return self.gcs_client.upload(request)
+
+    def check_nudity(self, auth: Authentication, request: NudityCheckRequest) -> NudityCheckResult:
+        """Check for nudity."""
+        self.accounts.authenticate(auth)
+
+        if not self.nudity_check_client:
+            self.nudity_check_client = self.client.get_nudity_check_client()
+
+        assert self.nudity_check_client
+        return self.nudity_check_client.check_nudity(request)
+
+    def compute_clip_score(self, auth: Authentication, request: CLIPScoreRequest) -> CLIPScoreResult:
+        """Computes CLIPScore for a given caption and image."""
+        self.accounts.authenticate(auth)
+
+        if not self.clip_score_client:
+            self.clip_score_client = self.client.get_clip_score_client()
+
+        assert self.clip_score_client
+        return self.clip_score_client.compute_score(request)
+
     def get_toxicity_scores(self, auth: Authentication, request: PerspectiveAPIRequest) -> PerspectiveAPIRequestResult:
         @retry_request
         def get_toxicity_scores_with_retry(request: PerspectiveAPIRequest) -> PerspectiveAPIRequestResult:
@@ -133,6 +200,16 @@ class ServerService(Service):
         self.accounts.authenticate(auth)
         return get_toxicity_scores_with_retry(request)
 
+    def get_moderation_results(self, auth: Authentication, request: ModerationAPIRequest) -> ModerationAPIRequestResult:
+        @retry_request
+        def get_moderation_results_with_retry(request: ModerationAPIRequest) -> ModerationAPIRequestResult:
+            if not self.moderation_api_client:
+                self.moderation_api_client = self.client.get_moderation_api_client()
+            return self.moderation_api_client.get_moderation_results(request)
+
+        self.accounts.authenticate(auth)
+        return get_moderation_results_with_retry(request)
+
     def make_critique_request(self, auth: Authentication, request: CritiqueRequest) -> CritiqueRequestResult:
         self.accounts.authenticate(auth)
         return self.client.get_critique_client().make_critique_request(request)
@@ -168,3 +245,6 @@ class ServerService(Service):
         hlog(f"Shutting down server by killing its own process {pid}...")
         os.kill(pid, signal.SIGTERM)
         hlog("Done.")
+
+    def get_cache_config(self, shard_name: str) -> CacheConfig:
+        return self.cache_backend_config.get_cache_config(shard_name)

helm/proxy/services/service.py

@@ -5,7 +5,11 @@ from typing import Dict, List, Tuple, Any
 
 from helm.common.general import parse_hocon
 from helm.common.critique_request import CritiqueRequest, CritiqueRequestResult
+from helm.common.clip_score_request import CLIPScoreRequest, CLIPScoreResult
+from helm.common.file_upload_request import FileUploadResult, FileUploadRequest
+from helm.common.nudity_check_request import NudityCheckRequest, NudityCheckResult
 from helm.common.perspective_api_request import PerspectiveAPIRequestResult, PerspectiveAPIRequest
+from helm.common.moderations_api_request import ModerationAPIRequest, ModerationAPIRequestResult
 from helm.common.tokenization_request import (
     WindowServiceInfo,
     TokenizationRequest,
@@ -17,6 +21,7 @@ from helm.common.request import Request, RequestResult
 from helm.benchmark.model_metadata_registry import ModelMetadata
 from helm.proxy.query import Query, QueryResult
 from helm.proxy.accounts import Authentication, Account
+from helm.common.cache import CacheConfig
 
 VERSION = "1.0"
 ACCOUNTS_FILE = "accounts.sqlite"
@@ -105,11 +110,31 @@ class Service(ABC):
         """Decodes to text."""
         pass
 
+    @abstractmethod
+    def upload(self, auth: Authentication, request: FileUploadRequest) -> FileUploadResult:
+        """Uploads a file to external storage."""
+        pass
+
+    @abstractmethod
+    def check_nudity(self, auth: Authentication, request: NudityCheckRequest) -> NudityCheckResult:
+        """Check for nudity for a batch of images."""
+        pass
+
+    @abstractmethod
+    def compute_clip_score(self, auth: Authentication, request: CLIPScoreRequest) -> CLIPScoreResult:
+        """Computes CLIPScore for a given caption and image."""
+        pass
+
     @abstractmethod
     def get_toxicity_scores(self, auth: Authentication, request: PerspectiveAPIRequest) -> PerspectiveAPIRequestResult:
         """Get toxicity scores for a batch of text."""
         pass
 
+    @abstractmethod
+    def get_moderation_results(self, auth: Authentication, request: ModerationAPIRequest) -> ModerationAPIRequestResult:
+        """Get OpenAI's moderation results for some text."""
+        pass
+
     @abstractmethod
     def make_critique_request(self, auth: Authentication, request: CritiqueRequest) -> CritiqueRequestResult:
         """Get responses to a critique request."""
@@ -149,3 +174,8 @@ class Service(ABC):
     def shutdown(self, auth: Authentication):
         """Shutdown server."""
         pass
+
+    @abstractmethod
+    def get_cache_config(self, shard_name: str) -> CacheConfig:
+        """Returns a CacheConfig"""
+        pass

helm/proxy/services/test_remote_service.py

@@ -17,7 +17,7 @@ from sqlitedict import SqliteDict
 from helm.common.authentication import Authentication
 from helm.common.request import Request, RequestResult
 from helm.common.tokenization_request import TokenizationRequest, TokenizationRequestResult
-from helm.proxy.accounts import Account
+from helm.proxy.accounts import Account, set_default_quotas
 from .remote_service import RemoteService
 from .service import ACCOUNTS_FILE
 
@@ -55,6 +55,7 @@ class TestRemoteServerService:
 
             with SqliteDict(os.path.join(path, ACCOUNTS_FILE)) as cache:
                 account: Account = Account(TestRemoteServerService._ADMIN_API_KEY, is_admin=True)
+                set_default_quotas(account)
                 cache[TestRemoteServerService._ADMIN_API_KEY] = asdict(account)
                 cache.commit()
             return path
@@ -126,9 +127,9 @@ class TestRemoteServerService:
         assert response.success
 
     def test_tokenize(self):
-        request = TokenizationRequest(text="1 2 3", tokenizer="simple/model1")
+        request = TokenizationRequest(text="1 2 3", tokenizer="simple/tokenizer1")
         response: TokenizationRequestResult = self.service.tokenize(self.auth, request)
-        assert [token.value for token in response.tokens] == ["1", "2", "3"]
+        assert [token.value for token in response.tokens] == ["1", " ", "2", " ", "3"]
 
     def test_make_request_plus_sign(self):
         # Ensure + in prompt doesn't get replaced by a blank space

helm/proxy/services/test_service.py

@@ -197,18 +197,6 @@ def helper_prod_test_service(request: Request, expected_text: str):
         # Consistency of log probs
         assert completion.logprob == sum(token.logprob for token in completion.tokens)
 
-        for token in completion.tokens[1:]:
-            assert len(token.top_logprobs) == request.top_k_per_token
-
-            # If generated token was one of the top, make sure has the right probability
-            if token.text in token.top_logprobs:
-                assert token.logprob == token.top_logprobs[token.text]
-
-            # If temperature = 0, then make sure we're getting the top probability token
-            if request.temperature == 0:
-                assert token.text in token.top_logprobs
-                assert token.logprob == max(token.top_logprobs.values())
-
     # Make sure we get the expected_text in one of the completions
     assert any(completion.text == expected_text for completion in result.completions)
 

helm/proxy/test_accounts.py

@@ -0,0 +1,32 @@
+import os
+import pytest
+import tempfile
+
+from helm.proxy.accounts import Accounts, Authentication, InsufficientQuotaError, Usage
+
+
+class TestAutoTokenCounter:
+    def setup_method(self, method):
+        accounts_file = tempfile.NamedTemporaryFile(delete=False)
+        self.accounts_path: str = accounts_file.name
+        self.accounts = Accounts(self.accounts_path)
+        self.root_auth = Authentication(Accounts.DEFAULT_API_KEY)
+
+    def teardown_method(self, method):
+        os.remove(self.accounts_path)
+
+    def test_check_can_use(self):
+        model_group = "anthropic"
+        account = self.accounts.create_account(self.root_auth)
+
+        # Cannot use this account because no quota was added
+        with pytest.raises(InsufficientQuotaError):
+            self.accounts.check_can_use(account.api_key, model_group)
+
+        # Add monthly quota
+        account.usages[model_group] = {}
+        account.usages[model_group]["monthly"] = Usage(quota=1000)
+        self.accounts.update_account(self.root_auth, account)
+
+        # Now this account has quota and can be used
+        self.accounts.check_can_use(account.api_key, model_group)