From 5a4ce9ef2b9f5e003ec5bebfa3647be2c4cb83db Mon Sep 17 00:00:00 2001 From: Christoph Grotz Date: Fri, 8 Sep 2023 18:12:24 +0200 Subject: [PATCH] VertexAI now allows to tune codey models (#10367) Description: VertexAI now supports to tune codey models, I adapted the Vertex AI LLM wrapper accordingly https://cloud.google.com/vertex-ai/docs/generative-ai/models/tune-code-models --- libs/langchain/langchain/llms/vertexai.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/libs/langchain/langchain/llms/vertexai.py b/libs/langchain/langchain/llms/vertexai.py index 2426587f58d..aaa5efbecb7 100644 --- a/libs/langchain/langchain/llms/vertexai.py +++ b/libs/langchain/langchain/llms/vertexai.py @@ -169,7 +169,7 @@ class VertexAI(_VertexAICommon, LLM): tuned_model_name = values.get("tuned_model_name") model_name = values["model_name"] try: - if tuned_model_name or not is_codey_model(model_name): + if not is_codey_model(model_name): from vertexai.preview.language_models import TextGenerationModel if tuned_model_name: @@ -181,7 +181,12 @@ class VertexAI(_VertexAICommon, LLM): else: from vertexai.preview.language_models import CodeGenerationModel - values["client"] = CodeGenerationModel.from_pretrained(model_name) + if tuned_model_name: + values["client"] = CodeGenerationModel.get_tuned_model( + tuned_model_name + ) + else: + values["client"] = CodeGenerationModel.from_pretrained(model_name) except ImportError: raise_vertex_import_error() return values