165 lines
6.1 KiB
Python
165 lines
6.1 KiB
Python
|
# -*- coding: utf-8 -*-
|
||
|
|
# Copyright 2023 Google LLC
|
||
|
|
#
|
||
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
|
# you may not use this file except in compliance with the License.
|
||
|
|
# You may obtain a copy of the License at
|
||
|
|
#
|
||
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
#
|
||
|
|
# Unless required by applicable law or agreed to in writing, software
|
||
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
|
# See the License for the specific language governing permissions and
|
||
|
|
# limitations under the License.
|
||
|
|
"""Utilities for Commands.
|
||
|
|
|
||
|
|
Common methods for Commands such as RunCommand and CompileCommand.
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from typing import AbstractSet, Any, Callable, Sequence
|
||
|
|
|
||
|
|
from google.generativeai.notebook import ipython_env
|
||
|
|
from google.generativeai.notebook import model_registry
|
||
|
|
from google.generativeai.notebook import parsed_args_lib
|
||
|
|
from google.generativeai.notebook import post_process_utils
|
||
|
|
from google.generativeai.notebook.lib import llm_function
|
||
|
|
from google.generativeai.notebook.lib import llmfn_input_utils
|
||
|
|
from google.generativeai.notebook.lib import llmfn_output_row
|
||
|
|
from google.generativeai.notebook.lib import llmfn_outputs
|
||
|
|
from google.generativeai.notebook.lib import unique_fn
|
||
|
|
|
||
|
|
|
||
|
|
class _GroundTruthLLMFunction(llm_function.LLMFunction):
|
||
|
|
"""LLMFunction that returns pre-generated ground truth data."""
|
||
|
|
|
||
|
|
def __init__(self, data: Sequence[str]):
|
||
|
|
super().__init__(outputs_ipython_display_fn=None)
|
||
|
|
self._data = data
|
||
|
|
|
||
|
|
def get_placeholders(self) -> AbstractSet[str]:
|
||
|
|
# Ground truth is fixed and thus has no placeholders.
|
||
|
|
return frozenset({})
|
||
|
|
|
||
|
|
def _call_impl(
|
||
|
|
self, inputs: llmfn_input_utils.LLMFunctionInputs | None
|
||
|
|
) -> Sequence[llmfn_outputs.LLMFnOutputEntry]:
|
||
|
|
normalized_inputs = llmfn_input_utils.to_normalized_inputs(inputs)
|
||
|
|
if len(self._data) != len(normalized_inputs):
|
||
|
|
raise RuntimeError(
|
||
|
|
"Ground truth should have same number of entries as inputs: {} vs {}".format(
|
||
|
|
len(self._data), len(normalized_inputs)
|
||
|
|
)
|
||
|
|
)
|
||
|
|
|
||
|
|
outputs: list[llmfn_outputs.LLMFnOutputEntry] = []
|
||
|
|
for idx, (value, prompt_vars) in enumerate(zip(self._data, normalized_inputs)):
|
||
|
|
output_row = llmfn_output_row.LLMFnOutputRow(
|
||
|
|
data={
|
||
|
|
llmfn_outputs.ColumnNames.RESULT_NUM: 0,
|
||
|
|
llmfn_outputs.ColumnNames.TEXT_RESULT: value,
|
||
|
|
},
|
||
|
|
result_type=str,
|
||
|
|
)
|
||
|
|
outputs.append(
|
||
|
|
llmfn_outputs.LLMFnOutputEntry(
|
||
|
|
prompt_num=0,
|
||
|
|
input_num=idx,
|
||
|
|
prompt_vars=prompt_vars,
|
||
|
|
output_rows=[output_row],
|
||
|
|
)
|
||
|
|
)
|
||
|
|
return outputs
|
||
|
|
|
||
|
|
|
||
|
|
def _get_ipython_display_fn(
|
||
|
|
env: ipython_env.IPythonEnv,
|
||
|
|
) -> Callable[[llmfn_outputs.LLMFnOutputs], None]:
|
||
|
|
return lambda x: env.display(x.as_pandas_dataframe())
|
||
|
|
|
||
|
|
|
||
|
|
def create_llm_function(
|
||
|
|
models: model_registry.ModelRegistry,
|
||
|
|
env: ipython_env.IPythonEnv | None,
|
||
|
|
parsed_args: parsed_args_lib.ParsedArgs,
|
||
|
|
cell_content: str,
|
||
|
|
post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
|
||
|
|
) -> llm_function.LLMFunction:
|
||
|
|
"""Creates an LLMFunction from Command.execute() arguments."""
|
||
|
|
prompts: list[str] = [cell_content]
|
||
|
|
|
||
|
|
llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None
|
||
|
|
|
||
|
|
llm_fn = llm_function.LLMFunctionImpl(
|
||
|
|
model=models.get_model(parsed_args.model_type),
|
||
|
|
model_args=parsed_args.model_args,
|
||
|
|
prompts=prompts,
|
||
|
|
outputs_ipython_display_fn=llmfn_outputs_display_fn,
|
||
|
|
)
|
||
|
|
if parsed_args.unique:
|
||
|
|
llm_fn = llm_fn.add_post_process_reorder_fn(name="unique", fn=unique_fn.unique_fn)
|
||
|
|
for fn in post_processing_fns:
|
||
|
|
llm_fn = fn.add_to_llm_function(llm_fn)
|
||
|
|
|
||
|
|
return llm_fn
|
||
|
|
|
||
|
|
|
||
|
|
def _convert_simple_compare_fn(
|
||
|
|
name_and_simple_fn: tuple[str, Callable[[str, str], Any]],
|
||
|
|
) -> tuple[str, llm_function.CompareFn]:
|
||
|
|
simple_fn = name_and_simple_fn[1]
|
||
|
|
new_fn = lambda x, y: simple_fn(x.result_value(), y.result_value())
|
||
|
|
return name_and_simple_fn[0], new_fn
|
||
|
|
|
||
|
|
|
||
|
|
def create_llm_compare_function(
|
||
|
|
env: ipython_env.IPythonEnv | None,
|
||
|
|
parsed_args: parsed_args_lib.ParsedArgs,
|
||
|
|
post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
|
||
|
|
) -> llm_function.LLMFunction:
|
||
|
|
"""Creates an LLMCompareFunction from Command.execute() arguments."""
|
||
|
|
llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None
|
||
|
|
|
||
|
|
llm_cmp_fn = llm_function.LLMCompareFunction(
|
||
|
|
lhs_name_and_fn=parsed_args.lhs_name_and_fn,
|
||
|
|
rhs_name_and_fn=parsed_args.rhs_name_and_fn,
|
||
|
|
compare_name_and_fns=[_convert_simple_compare_fn(x) for x in parsed_args.compare_fn],
|
||
|
|
outputs_ipython_display_fn=llmfn_outputs_display_fn,
|
||
|
|
)
|
||
|
|
for fn in post_processing_fns:
|
||
|
|
llm_cmp_fn = fn.add_to_llm_function(llm_cmp_fn)
|
||
|
|
|
||
|
|
return llm_cmp_fn
|
||
|
|
|
||
|
|
|
||
|
|
def create_llm_eval_function(
|
||
|
|
models: model_registry.ModelRegistry,
|
||
|
|
env: ipython_env.IPythonEnv | None,
|
||
|
|
parsed_args: parsed_args_lib.ParsedArgs,
|
||
|
|
cell_content: str,
|
||
|
|
post_processing_fns: Sequence[post_process_utils.ParsedPostProcessExpr],
|
||
|
|
) -> llm_function.LLMFunction:
|
||
|
|
"""Creates an LLMCompareFunction from Command.execute() arguments."""
|
||
|
|
llmfn_outputs_display_fn = _get_ipython_display_fn(env) if env else None
|
||
|
|
|
||
|
|
# First construct a regular LLMFunction from the cell contents.
|
||
|
|
llm_fn = create_llm_function(
|
||
|
|
models=models,
|
||
|
|
env=env,
|
||
|
|
parsed_args=parsed_args,
|
||
|
|
cell_content=cell_content,
|
||
|
|
post_processing_fns=post_processing_fns,
|
||
|
|
)
|
||
|
|
|
||
|
|
# Next create a LLMCompareFunction.
|
||
|
|
ground_truth_fn = _GroundTruthLLMFunction(data=parsed_args.ground_truth)
|
||
|
|
llm_cmp_fn = llm_function.LLMCompareFunction(
|
||
|
|
lhs_name_and_fn=("actual", llm_fn),
|
||
|
|
rhs_name_and_fn=("ground_truth", ground_truth_fn),
|
||
|
|
compare_name_and_fns=[_convert_simple_compare_fn(x) for x in parsed_args.compare_fn],
|
||
|
|
outputs_ipython_display_fn=llmfn_outputs_display_fn,
|
||
|
|
)
|
||
|
|
|
||
|
|
return llm_cmp_fn
|