Pencarian Film Menggunakan Milvus dan SentenceTransformers
Dalam contoh ini, kita akan mencari ringkasan plot film menggunakan Milvus dan pustaka SentenceTransformers. Dataset yang akan kita gunakan adalah Wikipedia Movie Plots with Summaries yang dihosting di HuggingFace.
Mari kita mulai!
Pustaka yang Dibutuhkan
Untuk contoh ini, kita akan menggunakan pymilvus
untuk terhubung menggunakan Milvus, sentence-transformers
untuk menghasilkan penyematan vektor, dan datasets
untuk mengunduh dataset contoh.
pip install pymilvus sentence-transformers datasets tqdm
from datasets import load_dataset
from pymilvus import MilvusClient
from pymilvus import FieldSchema, CollectionSchema, DataType
from sentence_transformers import SentenceTransformer
from tqdm import tqdm
Kita akan mendefinisikan beberapa parameter global,
embedding_dim = 384
collection_name = "movie_embeddings"
Mengunduh dan Membuka Dataset
Dalam satu baris, datasets
memungkinkan kita untuk mengunduh dan membuka dataset. Perpustakaan akan menyimpan dataset secara lokal dan menggunakan salinan tersebut pada saat dataset dijalankan. Setiap baris berisi rincian film yang memiliki artikel Wikipedia yang menyertainya. Kami menggunakan kolom Title
, PlotSummary
, Release Year
, dan Origin/Ethnicity
.
ds = load_dataset("vishnupriyavr/wiki-movie-plots-with-summaries", split="train")
print(ds)
Menghubungkan ke Basis Data
Pada titik ini, kita akan mulai menyiapkan Milvus. Langkah-langkahnya adalah sebagai berikut:
- Buatlah sebuah database Milvus Lite dalam sebuah berkas lokal. (Ganti URI ini dengan alamat server untuk Milvus Standalone dan Milvus Distributed).
client = MilvusClient(uri="./sentence_transformers_example.db")
- Buat skema data. Ini menentukan bidang-bidang yang terdiri dari sebuah elemen termasuk dimensi penyematan vektor.
fields = [
FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
FieldSchema(name="title", dtype=DataType.VARCHAR, max_length=256),
FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=embedding_dim),
FieldSchema(name="year", dtype=DataType.INT64),
FieldSchema(name="origin", dtype=DataType.VARCHAR, max_length=64),
]
schema = CollectionSchema(fields=fields, enable_dynamic_field=False)
client.create_collection(collection_name=collection_name, schema=schema)
- Tentukan algoritme pengindeksan pencarian vektor. Milvus Lite mendukung tipe indeks FLAT, sedangkan Milvus Standalone dan Milvus Distributed mengimplementasikan berbagai macam metode seperti IVF, HNSW, dan DiskANN. Untuk skala data yang kecil dalam demo ini, semua jenis indeks pencarian sudah cukup, jadi kami menggunakan yang paling sederhana, yaitu FLAT.
index_params = client.prepare_index_params()
index_params.add_index(field_name="embedding", index_type="FLAT", metric_type="IP")
client.create_index(collection_name, index_params)
Setelah langkah-langkah ini selesai, kita siap memasukkan data ke dalam koleksi dan melakukan pencarian. Setiap data yang ditambahkan akan diindeks secara otomatis dan langsung tersedia untuk pencarian. Jika datanya masih sangat baru, pencarian mungkin akan lebih lambat karena pencarian brute force akan digunakan pada data yang masih dalam proses pengindeksan.
Memasukkan Data
Untuk contoh ini, kita akan menggunakan model miniLM SentenceTransformers untuk membuat penyisipan teks plot. Model ini menghasilkan penyematan 384 dimensi.
model = SentenceTransformer("all-MiniLM-L12-v2")
Kami mengulang baris data, menyematkan bidang ringkasan plot, dan menyisipkan entitas ke dalam basis data vektor. Secara umum, Anda harus melakukan langkah ini pada sekumpulan item data untuk memaksimalkan throughput CPU atau GPU untuk model penyematan, seperti yang kita lakukan di sini.
for batch in tqdm(ds.batch(batch_size=512)):
embeddings = model.encode(batch["PlotSummary"])
data = [
{"title": title, "embedding": embedding, "year": year, "origin": origin}
for title, embedding, year, origin in zip(
batch["Title"], embeddings, batch["Release Year"], batch["Origin/Ethnicity"]
)
]
res = client.insert(collection_name=collection_name, data=data)
Operasi di atas relatif memakan waktu karena penyematan membutuhkan waktu. Langkah ini membutuhkan waktu sekitar 2 menit menggunakan CPU pada MacBook Pro 2023 dan akan jauh lebih cepat dengan GPU khusus. Beristirahatlah dan nikmati secangkir kopi!
Melakukan Pencarian
Dengan semua data yang telah dimasukkan ke dalam Milvus, kita dapat mulai melakukan pencarian. Dalam contoh ini, kita akan mencari film berdasarkan ringkasan plot dari Wikipedia. Karena kita melakukan pencarian batch, waktu pencarian dibagi untuk semua pencarian film. (Dapatkah Anda menebak film apa yang ingin saya ambil berdasarkan teks deskripsi kueri?)
queries = [
'A shark terrorizes an LA beach.',
'An archaeologist searches for ancient artifacts while fighting Nazis.',
'Teenagers in detention learn about themselves.',
'A teenager fakes illness to get off school and have adventures with two friends.',
'A young couple with a kid look after a hotel during winter and the husband goes insane.',
'Four turtles fight bad guys.'
]
# Search the database based on input text
def embed_query(data):
vectors = model.encode(data)
return [x for x in vectors]
query_vectors = embed_query(queries)
res = client.search(
collection_name=collection_name,
data=query_vectors,
filter='origin == "American" and year > 1945 and year < 2000',
anns_field="embedding",
limit=3,
output_fields=["title"],
)
for idx, hits in enumerate(res):
print("Query:", queries[idx])
print("Results:")
for hit in hits:
print(hit["entity"].get("title"), "(", round(hit["distance"], 2), ")")
print()
Hasilnya adalah:
Query: An archaeologist searches for ancient artifacts while fighting Nazis.
Results:
Love Slaves of the Amazons ( 0.4 )
A Time to Love and a Time to Die ( 0.39 )
The Fifth Element ( 0.39 )
Query: Teenagers in detention learn about themselves.
Results:
The Breakfast Club ( 0.54 )
Up the Academy ( 0.46 )
Fame ( 0.43 )
Query: A teenager fakes illness to get off school and have adventures with two friends.
Results:
Ferris Bueller's Day Off ( 0.48 )
Fever Lake ( 0.47 )
Losin' It ( 0.39 )
Query: A young couple with a kid look after a hotel during winter and the husband goes insane.
Results:
The Shining ( 0.48 )
The Four Seasons ( 0.42 )
Highball ( 0.41 )
Query: Four turtles fight bad guys.
Results:
Teenage Mutant Ninja Turtles II: The Secret of the Ooze ( 0.47 )
Devil May Hare ( 0.43 )
Attack of the Giant Leeches ( 0.42 )