Resolving review comments

RyanMullins · RyanMullins · commit c778fe12da2c · 2024-10-15T12:36:03.000Z
diff --git a/lit_nlp/examples/gcp/README.md b/lit_nlp/examples/gcp/README.md
@@ -27,6 +27,6 @@ docker build -f ./lit_nlp/examples/gcp/Dockerfile -t lit-app:gcp-dev .
 ### Run GPT-2 in a Docker container
 
 ```shell
-# Runs GPT-2 in Keras on Tensorflow
+# Runs GPT-2 in Keras on PyTorch
 docker run --rm -p 5432:5432 -e MODEL_CONFIG=gpt2:gpt2_base_en lit-app:gcp-dev
 ```
diff --git a/lit_nlp/examples/gcp/model.py b/lit_nlp/examples/gcp/model.py
@@ -13,7 +13,7 @@
 # limitations under the License.
 # ==============================================================================
 
-"""Wrapper for connetecting to LLMs on GCP via the model_server HTTP API."""
+"""Wrapper for connecting to LLMs on GCP via the model_server HTTP API."""
 
 from lit_nlp import app as lit_app
 from lit_nlp.api import model as lit_model
@@ -24,11 +24,6 @@
 from lit_nlp.lib import serialize
 import requests
 
-"""
-Plan for this module:
-
-"""
-
 _LlmHTTPEndpoints = lit_gcp_constants.LlmHTTPEndpoints
 
 LLM_ON_GCP_INIT_SPEC: lit_types.Spec = {
diff --git a/lit_nlp/examples/gcp/model_server.py b/lit_nlp/examples/gcp/model_server.py
@@ -18,21 +18,14 @@
 from collections.abc import Sequence
 import functools
 import os
-from typing import Optional
 
 from absl import app
-from lit_nlp import dev_server
 from lit_nlp.examples.gcp import constants as lit_gcp_constants
 from lit_nlp.examples.prompt_debugging import models as pd_models
 from lit_nlp.examples.prompt_debugging import utils as pd_utils
 from lit_nlp.lib import serialize
 from lit_nlp.lib import wsgi_app
 
-DEFAULT_DL_FRAMEWORK = 'kerasnlp'
-DEFAULT_DL_RUNTIME = 'tensorflow'
-DEFAULT_PRECISION = 'bfloat16'
-DEFAULT_SEQUENCE_LENGTH = 512
-DEFAULT_BATCH_SIZE = 1
 DEFAULT_MODELS = 'gemma_1.1_2b_IT:gemma_1.1_instruct_2b_en'
 
 _LlmHTTPEndpoints = lit_gcp_constants.LlmHTTPEndpoints
@@ -59,11 +52,13 @@ def _handler(app: wsgi_app.App, request, unused_environ):
         f'Only 1 model configuration can be provided, got {num_configs}'
     )
 
-  dl_framework = os.getenv('DL_FRAMEWORK', DEFAULT_DL_FRAMEWORK)
-  dl_runtime = os.getenv('DL_RUNTIME', DEFAULT_DL_RUNTIME)
-  precision = os.getenv('PRECISION', DEFAULT_PRECISION)
-  batch_size = int(os.getenv('BATCH_SIZE', DEFAULT_BATCH_SIZE))
-  sequence_length = int(os.getenv('SEQUENCE_LENGTH', DEFAULT_SEQUENCE_LENGTH))
+  dl_framework = os.getenv('DL_FRAMEWORK', pd_models.DEFAULT_DL_FRAMEWORK)
+  dl_runtime = os.getenv('DL_RUNTIME', pd_models.DEFAULT_DL_RUNTIME)
+  precision = os.getenv('PRECISION', pd_models.DEFAULT_PRECISION)
+  batch_size = int(os.getenv('BATCH_SIZE', pd_models.DEFAULT_BATCH_SIZE))
+  sequence_length = int(
+      os.getenv('SEQUENCE_LENGTH', pd_models.DEFAULT_SEQUENCE_LENGTH)
+  )
 
   models = pd_models.get_models(
       models_config=model_config,
@@ -92,7 +87,8 @@ def _handler(app: wsgi_app.App, request, unused_environ):
       wrapped_handlers, project_root='gcp', index_file='index.html'
   )
 
-def main(argv: Sequence[str]) -> Optional[dev_server.LitServerType]:
+
+def main(argv: Sequence[str]) -> wsgi_app.App:
   if len(argv) > 1:
     raise app.UsageError('Too many command-line arguments.')
 
diff --git a/lit_nlp/examples/gcp/model_server_gunicorn_config.py b/lit_nlp/examples/gcp/model_server_gunicorn_config.py
@@ -22,4 +22,4 @@
 timeout = 3600
 threads = 8
 worker_class = 'gthread'
-wsgi_app = f'lit_nlp.examples.gcp.model_server:get_wsgi_app()'
+wsgi_app = 'lit_nlp.examples.gcp.model_server:get_wsgi_app()'
diff --git a/lit_nlp/examples/gcp/server_gunicorn_config.py b/lit_nlp/examples/gcp/server_gunicorn_config.py
@@ -22,4 +22,4 @@
 timeout = 3600
 threads = 8
 worker_class = 'gthread'
-wsgi_app = f'lit_nlp.examples.gcp.server:get_wsgi_app()'
+wsgi_app = 'lit_nlp.examples.gcp.server:get_wsgi_app()'
diff --git a/lit_nlp/examples/gcp/vertexai_models.py b/lit_nlp/examples/gcp/vertexai_models.py
@@ -25,15 +25,11 @@
 from lit_nlp.api import types as lit_types
 
 _MAX_NUM_RETRIES = 5
-
 _DEFAULT_CANDIDATE_COUNT = 1
-
 _DEFAULT_MAX_OUTPUT_TOKENS = 256
-
 _DEFAULT_TEMPERATURE = 0.7
 
 
-# TODO(fanyecourage): Rename this class.
 class GeminiFoundationalModel(lit_model.BatchedRemoteModel):
   """GeminiFoundationalModel is a wrapper for foundatoinal Gemini models in Vertex AI Model Garden model.
 
