利用Embedding优化搜索功能

我们继续用Gemini学习LLM编程之旅。

Embedding是一种自然语言处理 (NLP) 技术，可将文本转换为数值向量。Embedding捕获语义含义和上下文，从而导致具有相似含义的文本具有更接近的Embedding。例如，句子“我带我的狗去看兽医”和“我带我的猫去看兽医”在向量空间中的Embedding会比较接近，因为它们都描述了相似的上下文。

Gemini 接口中的嵌入服务（models/embedding-001）可为单词、短语和句子生成embedding，然后生成的embedding可用于 NLP 任务，例如语义搜索、文本分类和聚类等。

本期文章参考徐文浩老师的《语义检索，利用Embedding优化你的搜索功能》(https://time.geekbang.org/column/article/644795) ，学习如何用embedding优化搜索，即通过语义来搜索（在给定一段输入文本的情况下检索语义相似的文本），而非传统的基于关键词分词的搜索。徐文浩老师的课程中用OpenAI的text-davinci-003模型生成淘宝产品标题，用text-embedding-ada-002生成Embedding。

由于gemini-pro对中文支持不好，models/embedding-001甚至完全不支持中文，我们会用gemini-pro生成Amazon的英文产品标题，用models/embedding-001生成Embedding。

一开始没留意谷歌的models/embedding-001不支持中文（文档没提），只是奇怪计算出来的余弦相似度很多一样，还以为自己的代码错了。后面用代码具体测试，才发现它生成的中文embedding一模一样。

human = genai.embed_content(model="models/embedding-001", content='人类', task_type="retrieval_query",
)['embedding']
mouse = genai.embed_content(model="models/embedding-001", content='鼠标', task_type="retrieval_query",
)['embedding']

# 比较两个嵌入向量列表是否相同
embeddings_are_equal = human == mouse
print(f"人类和鼠标的嵌入向量是否相同: {embeddings_are_equal}")

得到的结果是：

人类和鼠标的嵌入向量是否相同: True

设置环境

先参考《免费使用谷歌Gemini模型学习LLM编程》（https://juejin.cn/post/7315009908302331923）获取API Key，设置环境变量。

import os

import google.generativeai as genai
import pandas as pd
from dotenv import find_dotenv, load_dotenv
import numpy as np

_ = load_dotenv(find_dotenv())  # read local .env file

GOOGLE_API_KEY = os.getenv('GOOGLE_API_KEY')
genai.configure(api_key=GOOGLE_API_KEY)

model = genai.GenerativeModel("gemini-pro")

通过Gemini Pro生成实验数据

一般可以从网上下载或者直接用机器学习软件包自带的数据集来做实验，现在大语言模型来了，我们也可以直接让大语言模型帮我们生成实验数据。只要不是用于训练大模型，直接使用没有任何问题。

（参考字节跳动账户被封禁）

data = model.generate_content(
    " Please generate 50 product titles from Amazon, each about 50 characters long. The category is 3C digital products. The title often contains some promotional information, one per line.").text
digital_product_names = data.strip().split('\n')
digital_df = pd.DataFrame({'product_name': digital_product_names})
digital_df.product_name = digital_df.product_name.apply(lambda x: x.split('.')[1].strip())
digital_df.head()

让gemini-pro生成amazon的3C类别的产品标题，按行分隔，并加载到panda的DataFrame。注意，大模型返回的结果里面带了标号，通过这行代码去掉了：digital_df.product_name = digital_df.product_name.apply(lambda x: x.split('.')[1].strip())

输出结果：

All-in-One 4K UHD Smart TV with Built-in Streaming Apps and Voice Control
High-Performance Gaming PC with RGB Lighting and Liquid Cooling
Ultra-Thin Laptop with Long Battery Life and Fast Charging
Premium Noise-Canceling Headphones with Hi-Res Audio and Multipoint Connectivity
Powerful Wireless Router with Mesh Technology for Whole-Home Coverage

为了让实验数据多样化，再让模型生成女士的衣服、箱包类的商品标题。

clothes_data = model.generate_content(
    """Please generate 50 product titles from Amazon, each about 50 characters long. The category is women’s clothing, bags, etc. The title often contains some promotional information, one per line.""").text
clothes_product_names = clothes_data.strip().split('\n')
clothes_df = pd.DataFrame({'product_name': clothes_product_names})
clothes_df.product_name = clothes_df.product_name.apply(lambda x: x.split('.')[1].strip())
clothes_df.head()

输出结果：

Women’s Lightweight Summer Maxi Dress with Pockets: Flattering Flowy Beach Style
Crossbody Bag for Women: Durable Leather with Adjustable Strap and Multiple Compartments
“Women’s Designer Handbag Set: Tote, Crossbody, and Clutch Bag in One”
Vegan Leather Backpack for Women: Convertible to Tote or Shoulder Bag
Women’s Casual Shift Dress with Ruffle Sleeves: Comfortable and Stylish

再把两个DataFrame拼接起来作为我们最终的实验数据：

df = pd.concat([digital_df, clothes_df], axis=0)
df = df.reset_index(drop=True)
df

输出结果：

通过 Embedding 进行语义搜索

传统的搜索主要有两种：一种是数据库，不管是关系型数据库还是NoSQL数据库，都提供like类型的搜索，这种方式是直接查找某个字段是否包含某个关键词；另外一种是像ElasticSearch之类的全文检索，通过倒排索引的方式，先建立关键词到具体文档的映射，在搜索的时候，匹配上关键词，再返回对应的文档列表。

这两种方式具体到细节，当然都有很多可以优化的地方，但是本质上，都是通过关键词匹配，你的关键词选错了，就找不到对应的商品、文档等。所以，很多电商的商家，在将产品上架的时候，会在系统允许的范围内将商品的标题填上各种各样的关键词，目的就是为了尽可能匹配上关键词，获得更多搜索的流量。后台也会提供类似“买家搜索词”之类的工具，方便商家了解用户可能用什么关键词来搜索。这种“人工”智能比较费人工，但是关键词堆多了，也能有比较好的效果。

这个策略最大的缺点就是如果产品没有包含那个关键词，即使是同义词，语义上很接近，也搜索不到。

现在有了Embedding 的接口，就可以把一段文本的语义表示成一段向量。而向量之间是可以计算距离的。我们可以把用户的搜索也通过 Embedding 接口变成向量，然后把它和所有的商品的标题向量计算一下点积或者余弦距离，找出离我们搜索词最近的几个向量。最近的几个向量，其实就是语义和这个商品相似的，而并不需要相同的关键词。

下面来看一下具体的步骤。

保存文档的Embedding

首先，我们要把随机生成的所有商品标题，都计算出它们的Embedding，然后保存下来。

product_names_list = df.product_name.tolist()
#Note: Specifying a title for RETRIEVAL_DOCUMENT provides better quality embeddings for retrieval.
embeddings = (
    genai.embed_content(model="models/embedding-001", content=product_names_list, task_type="retrieval_document", title="amazon product names"))[
    'embedding']
print(f"embeddings len:{len(embeddings)}")
df["embedding"] = embeddings
df.to_parquet("data/amazon_product_title.parquet", index=False)
df.to_csv("data/amazon_product.csv", index=False)
df

注意谷歌生成embedding和其他的不太一样，有一个task_type，可能是针对不同的应用场景做了优化。

这里我们是用来做文档检索，而且现在是建立文档的“数据库”，所以task_type设置为retrieval_document，针对retrieval_document，还可以设置一个可选的title；等下查询的时候task_type要设为retrieval_query。

谷歌的models/embedding-001生成embedding还是很方便的，传入字符串列表就可以直接生成向量列表。

Task Type	Description
RETRIEVAL_QUERY	Specifies the given text is a query in a search/retrieval setting.
RETRIEVAL_DOCUMENT	Specifies the given text is a document in a search/retrieval setting.
SEMANTIC_SIMILARITY	Specifies the given text will be used for Semantic Textual Similarity (STS).
CLASSIFICATION	Specifies that the embeddings will be used for classification.
CLUSTERING	Specifies that the embeddings will be used for clustering.

搜索产品

接下来我们定义一个 search_product 的搜索函数，接受三个参数：df 代表用于搜索的数据源， query 代表用于搜索的搜索词，n 代表搜索返回多少条记录。

而这个函数做了这三件事情：

• 调用谷歌的接口将搜索词转换成 Embedding。
• 将这个 Embedding 和 DataFrame 里的每一个 Embedding 都计算一下点积。
• 根据点积去排序，返回点积最大的 n 个标题。

点积的值可以在 -1 和 1 之间（包含 -1 和 1）。如果两个向量之间的点积为 1，则这两个向量的方向相同。如果点积值为 0，则这些向量彼此正交或不相关。最后，如果点积为 -1，则向量指向相反方向并且彼此不相似。

def search_product(df, query, n=5, pprint=True):
    product_embedding = genai.embed_content(model="models/embedding-001", content=query, task_type="retrieval_query")[
        'embedding']
    df["dot_products"] = df.embedding.apply(lambda x: np.dot(x, product_embedding))
    results = (
        df.sort_values("dot_products", ascending=False).head(n)
        .loc[:, ["product_name", "dot_products"]]
    )
    if pprint:
        for index, row in results.iterrows():
            print(f"{row['product_name']} - dot_products: {row['dot_products']}")
    return results

results = search_product(df, "elegant dress for summer for beach vacations", n=3)

我们搜索 elegant dress for summer for beach vacations 得到结果：

Women's Lightweight Summer Maxi Dress with Pockets: Flattering Flowy Beach Style - dot_products: 0.6860313467181551
Women's Mini Dress: Sexy and Elegant, Perfect for a Night Out - dot_products: 0.6560099266011095
Women's Floral Print Wrap Dress: Flattering Design for Any Occasion - dot_products: 0.6244976428038793

搜索英文的结果还可以。

注意事项

在实际项目中使用这种方式优化的时候，或者是使用向量数据库来找匹配的文档的时候，要注意的是，返回的文档只是库里和你的查询最相似的。但是即使相差十万八千里，向量的距离也不太可能是0或者-1。以我们这次实验的数据为例，“elegant dress for summer for beach vacations”和”VR Gaming Controller with Haptic Feedback and Motion Controls”是最不相关的，点积也有0.415979671305287。所以，可能要根据具体的项目，低于某个数值就认为是不相干的，就不要返回了。

参考

1. https://time.geekbang.org/column/article/644795
2. https://ai.google.dev/docs/embeddings_guide?hl=en
3. https://ai.google.dev/examples/doc_search_emb