１．はじめに
google vertex aiのtextembedding-gecko-multilingualを使って、各国特許の要約のベクトル表現を作成して検索しました。
２．ベクトルデータベース作りと検索
検索をする際にキーワードではなく、文章からベクトルデータを使って、そのベクトルデータを使って検索したいと思います。
商用DBにはAPIが色々ありそうですが、とりあえず気軽に呼び出すべく、自作してみました。
（１）要約データを取得する
これは毎度使わせてもらっているgoogle patents public databaseで取得します。
全データだと１億レコードくらいになってしまうので、2020年以降かつCN特許を除いて取得します。
※CNデータを含むと大変な数（2400万件くらい）になってしまったので、省略しました。省略後でだいたい220万件ほど。
SELECT
family_id,
ARRAY_AGG(DISTINCT application_number) as appnum,
ARRAY_AGG(DISTINCT filing_date) as filing_date,
ARRAY_AGG(DISTINCT publication_number) as pubnum,
SPLIT(STRING_AGG(DISTINCT abst.text,"|"),"|")[0] as abst
FROM
`patents-public-data.patents.publications_202307`,
UNNEST(abstract_localized) AS abst
WHERE CAST(filing_date AS INT64) > 20200000
AND NOT SUBSTR(application_number, 0,2) = "CN"
GROUP BY family_id
（２）embedding
OpenAIのada(1536次元）もありましたが、やや割安かつVertex AIの中に含まれるtextembedding-gecko-multilingualでベクトル表現（768次元）に変換します。
!pip3 install google-cloud-aiplatform
!pip3 install shapely==2.0.0
exitPROJECT_ID = "***" # @param {type:"string"}
LOCATION = "us-central1" # @param {type:"string"}
import vertexai
vertexai.init(project=PROJECT_ID, location=LOCATION)
from vertexai.preview.language_models import TextEmbeddingModel
model = TextEmbeddingModel.from_pretrained("textembedding-gecko-multilingual@latest")こちらのソースコードを借りつつ、20000件づつ処理してjsonファイルとして出力していきます。
#@title jsonで作成してGCSに入れる。
#全レコード数
BQ_NUM_ROWS = 2191093 # @param {type:"integer"}
import tempfile
from pathlib import Path
# Create temporary file to write embeddings to
embeddings_file_path = Path(tempfile.mkdtemp())
print(f"Embeddings directory: {embeddings_file_path}")
#========================================================
import gc
import json
BQ_NUM_ROWS = int(BQ_NUM_ROWS) #上記のBIGQUERYで確認・取得した更新・追加対象の行数
BQ_CHUNK_SIZE = 20000 #20000行ずつ処理／json出力
BQ_NUM_CHUNKS = math.ceil(BQ_NUM_ROWS / BQ_CHUNK_SIZE)
START_CHUNK = 0
# Create a rate limit of 300 requests per minute. Adjust this depending on your quota.
API_CALLS_PER_SECOND = 600 / 60
# According to the docs, each request can process 5 instances per request
ITEMS_PER_REQUEST = 5
# Loop through each generated dataframe, convert
for i, df in tqdm(
enumerate(
query_bigquery_chunks(
max_rows=BQ_NUM_ROWS, rows_per_chunk=BQ_CHUNK_SIZE, start_chunk=START_CHUNK
)
),
total=BQ_NUM_CHUNKS - START_CHUNK,
position=-1,
desc="Chunk of rows from BigQuery",
):
# Create a unique output file for each chunk
chunk_path = embeddings_file_path.joinpath(
f"{embeddings_file_path.stem}_{i+START_CHUNK}.json"
)
with open(chunk_path, "a") as f:
id_chunk = df.id
#print(df)
# Convert batch to embeddings
is_successful, question_chunk_embeddings = encode_text_to_embedding_batched(
sentences=df["description"].tolist(),
api_calls_per_second=API_CALLS_PER_SECOND,
batch_size=ITEMS_PER_REQUEST,
)
# Append to file
embeddings_formatted = [
json.dumps(
{
"id": str(id),
"embedding": [str(value) for value in embedding],
}
)
+ "\n"
for id, embedding in zip(id_chunk[is_successful], question_chunk_embeddings)
]
f.writelines(embeddings_formatted)
# Delete the DataFrame and any other large data structures
del df
gc.collect()できたjsonファイル（110個）をGCSにアップロードします。
! gsutil -m cp -r {embeddings_file_path}/* {remote_folder}（３）json経由でBigQueryに入れる
BQのschemaを作っておけば、以下のコマンドで全部BQに入れてくれます。idが最初のfamily_idに対応します。
!bq --location=US load --source_format=NEWLINE_DELIMITED_JSON {dataset}.{table} gs://***/20240121*.json schema.json
（４）pineconeでベクトルデータベースに投入
BigQueryの内積計算クエリで類似文章も取れますが、費用と速度から、pinecone serverlessにしました。pinecone serverlessは今年pineconeに追加されたサービスで、従来のpineconeのベクトルデータベースよりも速く安いのが売りです。
pineconeの画面でベクトルデータベースを作っておけば、下記のようにメインは数行で投入できます。
!pip install pinecone-client -q
#colaboratory用
from google.colab import userdata
PINECONE_API_KEY = userdata.get('PINECONE_API_KEY')
from pinecone import Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index('patent-abst')frameは、上記で作成したテーブルデータをそのまま取ってきたものです。
index.upsert_from_dataframe(frame)
なんか件数が減ってしまいましたが、後でUPSERTできなかったものは再実行しないとです。
３．試し
入力文章を(textembedding-gecko-multilingualで)ベクトル表現に変更して、そのベクトル表現でpineconeを検索します。ライブラリなどは上記のものをそのまま使えます。
import vertexai
vertexai.init(project=PROJECT_ID, location=LOCATION)
from vertexai.preview.language_models import TextEmbeddingModel
model = TextEmbeddingModel.from_pretrained("textembedding-gecko-multilingual@latest")
def encode_texts_to_embeddings(sentences: List[str]) -> List[Optional[List[float]]]:
try:
embeddings = model.get_embeddings(sentences)
return [embedding.values for embedding in embeddings]
except Exception:
return [None for _ in range(len(sentences))]上記の関数で文章をベクトル化します。リストで帰ってくるのでvec_list[0]がほしいベクトルです。
text_list = ["雨雲が近づくのを検知して傘を開く"]
vec_list = encode_texts_to_embeddings(text_list)こんな感じにでるので、これを使って検索します。
from pinecone import Pinecone
pc = Pinecone(api_key=PINECONE_API_KEY)
index = pc.Index('patent-abst')
index.query(vector=vec_list[0],top_k=10)数秒で検索結果が出ます。idはfamily_idです。
{'matches': [{'id': '83063341', 'score': 0.831743479, 'values': []},
{'id': '81855280', 'score': 0.818136454, 'values': []},
{'id': '74670712', 'score': 0.811068118, 'values': []},
{'id': '81324479', 'score': 0.810769737, 'values': []},
{'id': '83721406', 'score': 0.808151245, 'values': []},
{'id': '81845510', 'score': 0.805805, 'values': []},
{'id': '85415761', 'score': 0.798141241, 'values': []},
{'id': '82657703', 'score': 0.796586096, 'values': []},
{'id': '79908436', 'score': 0.796158314, 'values': []},
{'id': '85600662', 'score': 0.795632243, 'values': []}],
'namespace': '',
'usage': {'read_units': 10}}上位３位はこちら。①③がイメージしていたのに近いです。
①id(family_id):83063341：TWM628300U
②id(family_id):81855280：JP3237832U
③id(family_id):74670712：TWM604947U
４．最後に
多言語を吸収してくれるembeddingnモデルのおかげですが、結構よい感じにmltilingualな検索をしてくれるなと思いました。今回のトータルの費用は6000円（embedding3000円、pineconeへの登録で2000円）くらいだったので、隙を見て中国文献も入れようと思います。
vertexAIで、画像や動画なども含めて検索するようにできればと作成中です。
Comments