使用代码模型和 Vertex AI SDK

Vertex AI 中有三个生成式 AI 代码基础模型。三个模型分别是代码生成模型、代码补全模型和代码聊天模型。

  • 代码生成模型名称为 code-bison，其在 Vertex AI SDK 中的类为 CodeGenerationModel。

  • 代码补全模型名称为 code-gecko，其在 Vertex AI SDK 中的类与用于生成代码的类 CodeGenerationModel 相同。

  • 代码聊天模型名称为 codechat-bison，其在 Vertex AI SDK 中的类为 CodeChatModel。

以下主题介绍了如何使用这些类和 Vertex AI SDK 来执行一些与代码相关的常见生成式 AI 任务。

生成代码函数

代码生成基础模型的用例包括设计单元测试、编写函数和创建类。要生成代码，请使用用于创建代码补全的类 CodeGenerationModel。如需创建生成代码的解决方案，请传入代码生成模型的版本名称，例如 code-bison@002。如需详细了解代码生成基础模型，请参阅创建提示以生成代码和测试代码生成提示。

以下代码示例编写了一个函数，用于检测某一年是否为闰年。

from vertexai.language_models import CodeGenerationModel

code_generation_model = CodeGenerationModel.from_pretrained("code-bison@001")
print(code_generation_model.predict(
    prefix="Write a function that checks if a year is a leap year.",
    # The following parameters are optional:
    # max_output_tokens=1024,
    # temperature=0.0,
))

输出可能类似以下内容：

I will write a function to check if a year is a leap year.

**The function will take a year as input and return a boolean value**.

**The function will first check if the year is divisible by 4.**

**If the year is divisible by 4, the function will then check if the year is divisible by 100.**

**If the year is divisible by 100, the function will then check if the year is divisible by 400.**

**If the year is divisible by 400, the function will ...

生成代码补全代码

代码补全是指在编写时预测补全代码的代码。系统会在您输入时显示生成的预测。如果您要创建代码补全解决方案，请使用 CodeGenerationModel 类。此类与用于生成代码（例如函数）的类相同。如需生成在编写时预测补全代码的代码，请调用 CodeGenerationModel.from_pretrained 并传入代码补全模型版本的名称。如需详细了解代码补全基础模型，请参阅创建代码补全提示和测试代码补全提示。

以下示例代码使用最新的稳定 code-gecko 版本，以返回预测补全反转字符串的函数开始部分的代码。

from vertexai.language_models import CodeGenerationModel

code_completion_model = CodeGenerationModel.from_pretrained("code-gecko@001")
print(code_completion_model.predict(
    prefix="def reverse_string(s):",
    # Optional:
    suffix="    return reversed_str",
    max_output_tokens=64,
    # temperature=0.0,
))

输出类似于以下内容：

"""
:type s: str
:rtype: str
"""
reversed_str = ""
for i in range(len(s) - 1, -1, -1):
    reversed_str += s[i]

创建代码聊天

您可能想要创建一个专门关于代码的聊天会话。例如，您可能希望通过聊天了解编码语言或语法。如需使用 Vertex AI SDK 创建代码聊天会话，请在 CodeChatModel 的实例上使用 start_chat 方法。与文本聊天不同，使用 Vertex AI SDK 创建的代码聊天不使用 InputOutputTextPair 类。如需详细了解代码聊天基础模型，请参阅创建代码聊天提示和测试代码聊天提示。

以下代码使用代码聊天来请求有关如何编写函数的信息。

from vertexai.language_models import CodeChatModel

code_chat_model = CodeChatModel.from_pretrained("codechat-bison@002")
code_chat = code_chat_model.start_chat()
print(code_chat.send_message("Please help write a function that prints its own source code"))

输出可能类似以下内容：

Sure, here is a function that prints its own source code:

```
def print_source_code():
  """Prints the source code of this function."""

  # Get the source code of this function.
  source_code = inspect.getsource(print_source_code)

  # Print the source code.
  print(source_code)

```

This function works by first getting the source code of the function using the
`inspect.getsource()` function. Then, it prints the source code to the console.

流式传输代码模型响应

您可能希望在生成代码和代码聊天模型时从中接收响应。在生成响应时从基础模型接收响应称为流式传输。流式传输代码生成和代码聊天模型响应时，输出令牌会在生成时发送。如需流式传输代码生成，请使用 CodeGenerationModel.predict_streaming 方法。如需流式传输代码聊天，请使用 CodeChatModel.predict_streaming 方法。如需详细了解如何从基础模型进行流式传输，请参阅从生成式 AI 模型流式传输响应。

流代码生成

以下示例代码流式传输了用于检查年份是否为闰年的代码。它还会输出调用 from_pretrained 之前和之后的时间，以显示流式传输输出所花费的时间。

import datetime
from vertexai.language_models import CodeGenerationModel

code_generation_model = CodeGenerationModel.from_pretrained("code-bison@001")

print("Start: ", datetime.datetime.now())
for response in code_generation_model.predict_streaming(
    prefix="Write a function that checks if a year is a leap year.",
    # Optional:
    # max_output_tokens=1024,
    # temperature=0.0,
):
    print(datetime.datetime.now(), "|", response)
print("End: ", datetime.datetime.now())

响应可能类似以下内容：

Start:  YYYY-MM-DD 06:31:45.759810
YYYY-MM-DD 06:31:46.536173 | To check if a year is a leap year, we can use the following step
YYYY-MM-DD 06:31:46.611856 | s:

1. **Check if the year is divisible by 4.** If it is not, th
YYYY-MM-DD 06:31:46.667330 | en it is not a leap year.
2. **Check if the year is divisible by
YYYY-MM-DD 06:31:46.805714 |  100.** If it is, then it is not a leap year unless it is also d
YYYY-MM-DD 06:31:46.940925 | ivisible by 400.
3. **If the year is divisible by 4 but not by 1
YYYY-MM-DD 06:31:47.033529 | 00, then it is a leap year.**

For example, the year 2020 is div
YYYY-MM-DD 06:31:47.110856 | isible by
End:  YYYY-MM-DD 06:31:47.112951

流式传输代码聊天

以下示例代码会流式传输响应聊天机器人请求的代码聊天，以编写一个输出自己的源代码的函数。此外，代码示例还会输出调用 from_pretrained 之前和之后的时间，以显示流式传输输出所花费的时间。

import datetime
from vertexai.language_models import CodeChatModel

code_chat_model = CodeChatModel.from_pretrained("codechat-bison@001")
code_chat = chat_model.start_chat()

print("Start: ", datetime.datetime.now())
for response in code_chat.send_message_streaming(
    message="Please help write a function that prints its own source code",
    # Optional:
    max_output_tokens=1024,
):
    #print(datetime.datetime.now(), "|", response)
    print(">>>")
    print(response)
print("End: ", datetime.datetime.now())

响应可能类似以下内容：

```
Start:  YYYY-MM-DD 06:32:10.733415
>>>
```python
def print_source(func):
    with open(func.__file__, '
>>>
r') as f:
        source = f.read()
    print(source)
```
End:  YYYY-MM-DD 06:32:11.764794
```

后续步骤