How to use Google's VertexAI PaLM 2 with Weights & Biases

# setup your GCP project and install the python SDK
# pip install google-cloud-aiplatform
project_id = "wandb-growth"
zone = "us-central1"

# initialize the Vertex AI API
vertexai.init(project=project_id, location=location)

model = TextGenerationModel.from_pretrained("text-bison@001")
    response = model.predict(
        "Give me ten interview questions for the role of program manager.",
        temperature=0.9,
    )

def palm_call(
    prompt: str,
    temperature: float = 0.7,
    max_output_tokens: int = 256,
    top_p: float = 0.8,
    top_k: int = 40,
    project_id: str = project_id,
    location: str = zone,
) -> str:
    vertexai.init(project=project_id, location=location)
    parameters = {
        "temperature": temperature,  # Temperature controls the degree of randomness in token selection.
        "max_output_tokens": max_output_tokens,  # Token limit determines the maximum amount of text output.
        "top_p": top_p,  # Tokens are selected from most probable to least until the sum of their probabilities equals the top_p value.
        "top_k": top_k,  # A top_k of 1 means the selected token is the most probable among all tokens.
    }

    model = TextGenerationModel.from_pretrained("text-bison@001")
    response = model.predict(
        prompt,
        **parameters,
    )
    return response.text

# let's define some configuration parameters
config = dict(
    temperature = 1.0,
    max_output_tokens = 128,
    top_p = 0.8,
    top_k = 40,
)

table = wandb.Table(columns=["model", "time", "temperature", "max_output_tokens", "top_p", "top_k", "prompt", "response"])

# we iterate through the queries and call the model
# adding the results to a table
for q in tqdm(queries):
    t0 = time.perf_counter()
    res = palm_call(q, **config)
    table.add_data(
        "text-bison@001", 
        time.perf_counter() - t0, 
        config["temperature"], 
        config["max_output_tokens"], 
        config["top_p"], 
        config["top_k"], 
        q, 
        res)

from langchain.llms import VertexAI

os.environ["LANGCHAIN_WANDB_TRACING"] = "true"

# create a wandb run on our project
wandb.init(project="wandb-palm", job_type="generation")

# we define as the LLM to use
llm = VertexAI(model="text-bison@001", project="wandb-growth", location="us-central1")

# we define a chain with custom or out-of-the-shelf tools
tools = [Tool1(), Tool2()]
agent = initialize_agent(
    tools, 
    llm, 
    agent=AgentType.ZERO_SHOT_REACT_DESCRIPTION,
    handle_parsing_errors=True,
    verbose=True
)

# W&B will keep track of the traces for you!
agent.run(
    "Do something fancy with my tools!"
)