玩轉 Python 與 MongoDB_Day29_Vector Search

2023 iThome 鐵人賽

DAY 29

Software Development

玩轉 Python 與 MongoDB系列第 29 篇

15th鐵人賽 python3 mongodb pymongo 教學

熊熊工程師

2023-09-29 08:58:08

332 瀏覽

分享至

今天我們要來介紹 MongoDB 在 7.0 版本以後新推出的 Vector Search 向量搜尋功能該如何實作，這個功能似乎目前只在 Atlas 上有支援，如果你的資料庫是部署在本地的或許還要再等等，當然如果有說錯歡迎留言告訴我！

今天示範的流程大約如下：

爬取 PTT 八卦版文章內文並寫入
利用 Hugging Face 上所提供的開源模型計算內文的 embedding 並寫入
在 Atlas 上建立 Vector Search 的索引
使用 Python 搭配 Vector Search 索引進行搜尋

在開始前，先附上今天會使用到的 pydantic schema，方便我們進行資料操作

from bson.objectid import ObjectId
from pydantic import BaseModel, Field
from datetime import datetime
from typing import List


class Paragraph(BaseModel):
    url: str
    title: str = None
    author: str = None
    content: str = None
    board: str = None
    post_time: datetime = None
    created_time: datetime = None
    title_embedding: List[float] = None


class DBParagraph(Paragraph):
    id: ObjectId = Field(alias="_id")

    class Config:
        arbitrary_types_allowed = True

一、爬取 PTT 八卦版文章

下方的爬蟲程式碼會去爬取 PTT 八卦版的文章並回傳，可以看到這個是一個 function，我們會在 demo 裡面進行呼叫並寫入資料庫，在今天的 GitHub 當中會附上完整的程式碼。另外由於本次目標是介紹 mongodb，下方範例中的爬蟲僅用於參考，不會精修。

import requests
from schema import Paragraph
from bs4 import BeautifulSoup
from datetime import datetime


def get_ptt_gossiping_paragraph() -> Paragraph:
    headers = {"cookie": "over18=1"}
    url = "https://www.ptt.cc/bbs/Gossiping/index.html"

    for i in range(10):
        response = requests.get(url=url, headers=headers)
        paragraph_list_soup = BeautifulSoup(response.text, "lxml")
        previous_page = paragraph_list_soup.select_one(
            "#action-bar-container > div > div.btn-group.btn-group-paging > a:nth-child(2)"
        )

        for paragraph in paragraph_list_soup.find_all(name="div", attrs={"class": "r-ent"}):
            if paragraph.find("a"):
                paragraph_response = requests.get(
                    url=f"https://www.ptt.cc{paragraph.find('a').get('href')}",
                    headers=headers
                )
                paragraph_soup = BeautifulSoup(paragraph_response.text, "lxml")
                if post_time := paragraph_soup.select_one("#main-content > div:nth-child(4) > span.article-meta-value"):
                    post_time = datetime.strptime(post_time.get_text(), '%a %b %d %H:%M:%S %Y')

                content = paragraph_soup.find(
                    name="div", attrs={"id": "main-content"}).get_text()
                author = paragraph_soup.select_one(
                    "#main-content > div:nth-child(1) > span.article-meta-value"
                )
                title = paragraph_soup.select_one(
                    "#main-content > div:nth-child(3) > span.article-meta-value"
                )

                if content and title and author:
                    print(title.text)
                    yield Paragraph(
                        post_time=post_time,
                        content=content,
                        title=title.text,
                        author=author.text,
                        board="Gossiping",
                        created_time=datetime.now(),
                        url=f"https://www.ptt.cc{paragraph.find('a').get('href')}"
                    )

        url = f"https://www.ptt.cc{previous_page.get('href')}"

執行完畢後可以看到我們成功在資料庫當中插入許多資料

ptt 文章截圖

二、使用 Hugging Face 計算 Embedding

取得 TOKEN

在登入後點選右上角的圓圈並點選 Settings 選項
選擇 Access Token 選項並點選建立新 Token，如果你是新使用者會需要先去驗證 Email

撰寫呼叫 API 計算 embedding 的程式，此 function 會存放在 parsers.py

本次使用到的模型為 all-MiniLM-L6-v2，和官方 MongoDB 在 demo 的時候一樣，有興趣的人可以自行看看 Hugging Face 上的其他模型

import os
import requests


def generate_embedding(text: str) -> list[float]:
    embedding_url = "https://api-inference.huggingface.co/pipeline/feature-extraction/sentence-transformers/all-MiniLM-L6-v2"
    response = requests.post(
        embedding_url,
        headers={"Authorization": f"Bearer {os.getenv('HUGGING_FACE_TOKEN')}"},
        json={"inputs": text})

    if response.status_code != 200:
        raise ValueError(f"Request failed with status code {response.status_code}: {response.text}")

    return response.json()

回到 demo.py 上呼叫 function 進行 embedding 的計算，下方可以看到計算後每個 document 的 title_embedding 欄位被成功替換成一個 list

for paragraph in collection.find():
  paragraph = DBParagraph(**paragraph)
  embedding = generate_embedding(text=paragraph.title)
  collection.update_one(
      {"_id": paragraph.id},
      {"$set": {"title_embedding": embedding}}
  )

計算成功

三、在 Atlas 上建立 Vector Search 索引

到 Atlas 上找到你的資料庫，並點選 Search 標籤選項
點選 "Create Search Index" 選項
選擇 "Json Editor" 選項並點選 "Next"
接著按照下方步驟開始建立索引
- 首先點選左方欄位選擇資料表
- 給予本次建立的索引一個名稱
- 針對內容的部分記得抽換成自己的欄位名稱，本次使用 content_embedding
- 點選 "Next" 按鈕
下方附上範例格式，另外 "dimensions" 欄位設定為 384 是因為我們選用的模型每次就是會解出 384 個 embedding，而針對 "similarity" 以及 "type" 欄位的詳細說明可以參考這個網址
```
{
    "mappings": {
        "dynamic": true,
        "fields": {
            "title_embedding": {
                "dimensions": 384,
                "similarity": "dotProduct",
                "type": "knnVector"
            }
        }
    }
}
```
畫面會跳轉至預覽頁面，點選 "Create Search Index" 選項
畫面會跳轉至成功建立，點選 "close" 按鈕

四、使用 Python 搭配 Vector Search 索引進行搜尋

接著我們回到程式上，我們可以透過 aggregate 他配 $search 運算符號來進行索引的操作，下方附上範例，另外針對 knnBeta 的搜尋參數的詳細介紹，可以參考這個網址

query = "Bella"

results = collection.aggregate([
    {
        '$search': {
            "index": "title_embedding_index",
            "knnBeta": {
                "vector": generate_embedding(query),
                "k": 5,
                "path": "title_embedding"
            }
        }
    }
])

for tmp in results:
    tmp = DBParagraph(**tmp)
    print(f"文章ＩＤ：{str(tmp.id)}，文章標題：{tmp.title}")

可以看到下方的搜尋結果，成功針對標題進行搜尋，而且印出的文章標題都與 Bella 有關

搜尋結果