用TextGrad优化LLM提示

LIBRARY Nov 17, 2024

在我们利用检索增强生成 (RAG) 构建 AI 应用程序的过程中，我们面临的挑战之一是针对 AI 需要处理的不同情况优化提示。我们正在寻找一种优化管道，可以在检测到特定模式时自动调整提示。就在那时，我们了解了 DSPy 和 TextGrad。

探索 DSPy 和 TextGrad

最初，我们探索了 DSPy，它专注于提示编程，而无需将提示直接硬编码到应用程序中。然而，这种方法与我们的需求无关，因为我们已经在工作流程的不同步骤中使用预定义的提示。

然后我们的注意力转移到 TextGrad。这个概念立即引起了我们的兴趣，我们决定对其进行测试，以了解其提示优化的潜力。在这篇文章中，我将带你了解我们使用 TextGrad 的经验以及如何使用它来优化 AI 工作流程中的提示。

1、设置不使用 TextGrad 的初始示例

以下是使用 OpenAI 的 API 从自然语言提示生成 Python 代码的基本示例：

from openai import OpenAI

client = OpenAI()
system_instructions = "You are a helpful assistant that specializes in generating Python code from natural language input"
user_input = "Generate a list of squares for numbers 1 to n."

response = client.chat.completions.create(
    model="gpt-4o-mini",
    messages=[
        {"role": "system", "content": system_instructions},
        {"role": "user", "content": user_input}
    ]
)

print(response.choices[0].message.content)

输出如下：

You can generate a list of squares for numbers from 1 to `n` using a simple Python list comprehension. Here's a code snippet that does that:
```python
def generate_squares(n):
    squares = [i ** 2 for i in range(1, n + 1)]
    return squares
# Example usage
n = 10
squares_list = generate_squares(n)
print(squares_list)
```
In this code, `generate_squares` takes an integer `n` as input and returns a list of squares from 1 to `n`. You can change the value of `n` to generate squares for a different range.

虽然输出是正确的，但它包含额外的解释。这是 TextGrad 可以帮助优化提示以专注于生成所需的 Python 代码的地方。

2、TextGrad基本概念

TextGrad 是一个 Python 包，它提供了一个简单的接口来实现用于文本优化的 LLM“梯度”。它使用类似于 PyTorch 的概念，例如变量、损失函数和优化器。

首先，你需要安装以下软件包：

pip install textgrad python-Levenshtein

TextGrad 的一些核心概念介绍如下：

变量：TextGrad 的构建块。
引擎：用于前向和后向传递的 LLM。
损失函数：评估生成的响应。
优化器：使用文本梯度下降基于后向传递优化变量。

这些组件改编自 PyTorch 概念，如果你习惯使用机器学习框架，则 TextGrad 会让你感到熟悉。

3、构建TextGrad提示优化管道

让我们来看一个示例，其中我们优化了提示以从自然语言问题生成 Python 代码。

让我们从导入所需的软件包和模块开始。然后我们加载我们希望优化提示的数据。

import textgrad as tg
import os
import pandas as pd
import numpy as np
from tqdm import tqdm
import os
np.random.seed(12)

#Dataset
data = [
    {
      "input": "Write a function to check if a number is prime.",
      "output": "def is_prime(n):\n    if n <= 1:\n        return False\n    for i in range(2, int(n**0.5) + 1):\n        if n % i == 0:\n            return False\n    return True"
    },
    {
      "input": "Generate a list of squares for numbers 1 to n.",
      "output": "def squares(n):\n    squares = [x**2 for x in range(1, n)]\n    print(squares)"
    },
    {
      "input": "Write a function to check if a given string is a palindrome.",
      "output": "def is_palindrome(s):\n    s = s.replace(' ', '').lower()\n    return s == s[::-1]"
    },
    {
      "input": "Create a function that returns the first n numbers in the Fibonacci sequence.",
      "output": "def fibonacci(n):\n    sequence = [0, 1]\n    for i in range(2, n):\n        sequence.append(sequence[-1] + sequence[-2])\n    return sequence[:n]"
    },
    {
      "input": "Write a Python function to calculate the factorial of a number using recursion.",
      "output": "def factorial(n):\n    if n == 0:\n        return 1\n    return n * factorial(n - 1)"
    },
    {
      "input": "Write a function to find the maximum value in a list.",
      "output": "def find_max(lst):\n    return max(lst)"
    },
    {
      "input": "Create a function to count the number of vowels in a given string.",
      "output": "def count_vowels(s):\n    return sum(1 for char in s.lower() if char in 'aeiou')"
    },
    {
      "input": "Write a function that sorts a list of integers in ascending order.",
      "output": "def sort_list(lst):\n    return sorted(lst)"
    },
    {
      "input": "Generate a function that checks if a number is even or odd.",
      "output": "def is_even(n):\n    return n % 2 == 0"
    },
    {
      "input": "Write a function to reverse a string.",
      "output": "def reverse_string(s):\n    return s[::-1]"
    }
  ]

dataset = pd.DataFrame(data)


random_int = np.random.randint(0, len(dataset)-1)
train_set = dataset.iloc[:random_int]
test_set = dataset.iloc[random_int:]

print(random_int, len(train_set), len(test_set))

3.1 系统提示初始化

在这里，我们将系统提示初始化为变量，并将 required_grad 设置为 True，以通知 TextGrad 继续更新每一步向后移动一个梯度。

SYSTEM_PROMPT = """
You are an expert assistant that specialises in generates Python code from natural language input.
Only generate python code. Do not add any comments or explanations.
"""
system_prompt_var = tg.Variable(SYSTEM_PROMPT, requires_grad=True, role_description="system prompt for Python code generation")

3.2 配置 TextGrad 引擎和优化器

然后我们将初始化引擎进行测试，该引擎从问题和评估引擎生成 Python 代码，该引擎将用于分析文本差异以改进提示。它与检查参数梯度下降的标准 ML 优化器相同。

一般来说，选择后向引擎为 LLM 是个好主意，它可以帮助提供前向传递 LLM，以便更好地处理提示。

现在最重要的部分是损失函数。你需要找到一个损失函数，它可以通过与组真值进行比较或分析来告诉优化器完成中的文本差异。

llm_engine = tg.get_engine("gpt-4o-mini")
model = tg.BlackboxLLM(llm_engine, system_prompt=system_prompt_var)

eval_engine = tg.get_engine("gpt-4o")
tg.set_backward_engine(eval_engine, override=True)

optimizer = tg.TGD(engine=eval_engine, parameters=[system_prompt_var])

3.3 定义损失函数

这里我采用了 StringBasedFunction，但你可以从文档中提供的众多损失函数中挑选任何一个。我挑选了 Levenshtein 距离来计算生成的代码与事实之间的差异，该距离为“不匹配”到“完全匹配”生成 0-1 之间的值。

我们可以挑选任何其他评估函数来将生成的完成与事实进行比较。

import Levenshtein

def levenshtein_accuracy(original, generated):
    distance = Levenshtein.distance(original, generated)
    max_len = max(len(original), len(generated))
    return (max_len - distance) / max_len if max_len > 0 else 1.0

def parse_code_answer(answer: str):
    import re
    pattern = r'```python\s+(.*?)\s+```'
    matches = re.findall(pattern, answer, re.DOTALL | re.IGNORECASE)
    return matches[0] if matches else answer

def string_based_equality_fn(prediction, ground_truth_answer):
    return levenshtein_accuracy(parse_code_answer(str(prediction.value)), parse_code_answer(str(ground_truth_answer.value)))

eval_fn = tg.autograd.StringBasedFunction(string_based_equality_fn, function_purpose="Levenshtein distance evaluation")

3.4 训练和评估

现在让我们开始构建训练循环。如前所述，这与 PyTorch 训练循环相同。

训练循环：我已经完成了 3 个时期，使用训练集计算损失，向后传播损失以调整文本梯度，然后逐步优化器将梯度纳入参数中。对于每个训练集，计算测试集准确率。如果准确率下降，我们会回退到之前的最佳提示。

def test_set_accuracy(step=1):
    accuracy = []
    for k in tqdm(range(len(test_set)), total=len(test_set), position=0, desc=f"Step {step}"):
        question, answer = test_set.iloc[k]
        question = tg.Variable(question, requires_grad=False, role_description="question")
        answer = tg.Variable(answer, requires_grad=False, role_description="original code")
        response = model(question)
        acc = levenshtein_accuracy(answer.value, response.value)
        accuracy.append(acc)
    return np.mean(accuracy)

result = {"prompt": [SYSTEM_PROMPT], "accuracy": []}
acc = test_set_accuracy(0)
result["accuracy"].append(acc)
print("Test Set Accuracy Before Training:", acc)

for epoch in range(1, 4):
    for i in tqdm(range(len(train_set)), total=len(train_set), position=0, desc=f"Epoch {epoch}"):
        question, answer = train_set.iloc[i]
        question = tg.Variable(question, requires_grad=False, role_description="question")
        answer = tg.Variable(answer, requires_grad=False, role_description="original code")
        response = model(question)
        response.set_role_description("generated code")
        loss = loss_fn(inputs=[question, response, answer])
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    
    test_acc = test_set_accuracy(epoch)
    print("Accuracy:", test_acc)
    if test_acc > result["accuracy"][-1]:
        result["prompt"].append(system_prompt_var.value)
        result["accuracy"].append(test_acc)
    else:
        result["prompt"].append(result["prompt"][-1])
        result["accuracy"].append(result["accuracy"][-1])
        system_prompt_var.value = result["prompt"][-1]