Passing parameters to Llama 2 deployed in Vertex AI

Question

I deployed Llama 2 Chat 13B in Vertex AI from model garden. I can also run inference, however, I can not pass any parameter other than prompt and temperature such as max_output_token, top_p, top_k.

from typing import Dict, List, Union

from google.cloud import aiplatform
from google.protobuf import json_format
from google.protobuf.struct_pb2 import Value


def predict_custom_trained_model_sample(
    project: str,
    endpoint_id: str,
    instances: Union[Dict, List[Dict]],
    location: str = "us-central1",
):
    """
    `instances` can be either single instance of type dict or a list
    of instances.
    """
    api_endpoint = f"{location}-aiplatform.googleapis.com"
    # The AI Platform services require regional API endpoints.
    client_options = {"api_endpoint": api_endpoint}
    # Initialize client that will be used to create and send requests.
    # This client only needs to be created once, and can be reused for multiple requests.
    client = aiplatform.gapic.PredictionServiceClient(client_options=client_options)
    # The format of each instance should conform to the deployed model's prediction input schema.
    instances = instances if isinstance(instances, list) else [instances]
    instances = [
        json_format.ParseDict(instance_dict, Value()) for instance_dict in instances
    ]
    parameters_dict = {}
    parameters = json_format.ParseDict(parameters_dict, Value())
    endpoint = client.endpoint_path(
        project=project, location=location, endpoint=endpoint_id
    )
    response = client.predict(
        endpoint=endpoint, instances=instances, parameters=parameters
    )
    print("Response")
    print("Deployed Model ID:", response.deployed_model_id)
    # The predictions are a google.protobuf.Value representation of the model's predictions.
    predictions = response.predictions
    for prediction in predictions:
        print("prediction:", prediction)
query = "Who is Albert Einstein?"
predict_custom_trained_model_sample(
    project="X",
    endpoint_id="Y",
    location="us-east1",
    instances = [
      {
         "prompt": query, "temperature": 0.0
      },
   ]
)

If I pass anything other than prompt or temperature in the dictionary, I get an error.

_InactiveRpcError                         Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/google/api_core/grpc_helpers.py in error_remapped_callable(*args, **kwargs)
     71         try:
---> 72             return callable_(*args, **kwargs)
     73         except grpc.RpcError as exc:

6 frames
_InactiveRpcError: <_InactiveRpcError of RPC that terminated with:
    status = StatusCode.INTERNAL
    details = "Internal Server Error"
    debug_error_string = "UNKNOWN:Error received from peer ipv4:74.125.203.95:443 {grpc_message:"Internal Server Error", grpc_status:13, created_time:"2023-10-31T22:34:55.857971339+00:00"}"
>

The above exception was the direct cause of the following exception:

InternalServerError                       Traceback (most recent call last)
/usr/local/lib/python3.10/dist-packages/google/api_core/grpc_helpers.py in error_remapped_callable(*args, **kwargs)
     72             return callable_(*args, **kwargs)
     73         except grpc.RpcError as exc:
---> 74             raise exceptions.from_grpc_error(exc) from exc
     75 
     76     return error_remapped_callable

InternalServerError: 500 Internal Server Error