电影推荐系统教程：使用Python和协同过滤构建Movie_Recommendation.ipynb
技术背景
传统推荐系统依赖协同过滤或基于内容的标签匹配，难以理解语义。Gemini API提供强大的文本嵌入能力，可将电影元数据转换为高维向量，捕捉深层语义。Qdrant作为高性能向量搜索引擎，支持亿级向量的相似性检索。两者结合，可实现基于自然语言查询的智能推荐，适用于个性化内容推荐、语义搜索等场景。但需注意API成本、数据规模限制及向量维度选择。
核心原理
核心流程：1) 数据准备：从Kaggle加载TMDB电影数据集，清洗并提取关键字段（标题、概述、类型等）。2) 文本构造：将字段拼接为结构化文本，供嵌入模型使用。3) 嵌入生成：调用Gemini API的gemini-embedding-001模型，批量生成3072维向量。4) 向量索引：将向量及元数据存入Qdrant集合，采用余弦距离度量。5) 语义搜索：用户查询经同一模型嵌入，在Qdrant中检索最相似向量，返回推荐电影。
架构分析
系统分为数据层、嵌入层、索引层和查询层。数据层使用Pandas处理CSV，嵌入层通过Gemini API异步批量生成向量，索引层利用Qdrant内存模式存储向量与payload，查询层提供recommend_movies函数。关键设计：使用BATCH_SIZE控制API请求大小，QDRANT_BATCH_SIZE优化upsert性能，重试机制增强鲁棒性。生产环境可替换为Qdrant云服务，并采用异步流水线提升吞吐。
完整实战教程
##### 版权所有 2026 Google LLC.
# @title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
使用 Gemini API 和 Qdrant 构建电影推荐系统
注意： 此笔记本需要付费层级速率限制才能正常运行。（更多详情请参阅定价）。
概述
Gemini API 提供了一系列生成式 AI 模型的访问权限，用于生成内容和解决问题。这些模型经过设计和训练，可以同时处理文本和图像作为输入。
Qdrant 是一个开源的向量相似度搜索引擎，专为高效、可扩展的语义搜索而设计。它提供了一个简单而强大的 API，用于存储和搜索高维向量，支持带有元数据（payloads）的过滤，并且可以轻松集成到生产系统中。Qdrant 可以自托管，也可以通过其托管云服务访问，这使得它能够快速设置，非常适合依赖语义理解和检索的各种 AI 应用。
在本笔记本中，您将学习如何借助 Gemini API 和 Qdrant 对网站数据进行相似度搜索。
<!-- 社区贡献者徽章 -->
<table>
<tr>
<!-- 作者头像单元格 -->
<td bgcolor="#d7e6ff">
<a href="https://github.com/andycandy" target="_blank" title="在 GitHub 上查看 Anand Roy 的个人资料">
</a>
</td>
<!-- 文本内容单元格 -->
<td bgcolor="#d7e6ff">
<h2><font color='black'>本笔记本由 <a href="https://github.com/andycandy" target="_blank"><font color='#217bfe'><strong>Anand Roy</strong></font></a> 贡献。</font></h2>
<h5><font color='black'><a href="https://www.linkedin.com/in/anand-roy-61a2b529b"><font color="#078efb">LinkedIn</font></a> - 查看 <a href="https://github.com/andycandy" target="_blank"><font color="#078efb"><strong>Anand</strong></font></a> 的其他笔记本 <a href="https://github.com/search?q=repo%3Agoogle-gemini%2Fcookbook%20%22Anand%20Roy%22&type=code" target="_blank"><font color="#078efb">此处</font></a>。</h5></font><br>
<!-- 页脚 -->
<font color='black'><small><em>有很酷的 Gemini 示例？欢迎 <a href="https://github.com/google-gemini/cookbook/blob/main/CONTRIBUTING.md" target="_blank"><font color="#078efb">也分享出来</font></a>！</em></small></font>
</td>
</tr>
</table>
环境设置
首先，您必须安装所需的软件包并设置必要的环境变量。
安装
安装 Google 的 Gemini API Python 客户端 SDK，即 google-genai。接下来，安装 Qdrant 的 Python 客户端 SDK，即 qdrant-client。
%pip install -q "google-genai>=2.9.0"
%pip install -q protobuf==4.25.1 qdrant-client[fastembed]
配置你的 API 密钥
要运行以下单元格，你的 API 密钥必须存储在名为 GEMINI_API_KEY 的 Colab 密钥中。如果你还没有 API 密钥，或者不确定如何创建 Colab 密钥，请参阅 身份验证 示例。
from google.colab import userdata
from google import genai
GEMINI_API_KEY=userdata.get('GEMINI_API_KEY')
client = genai.Client(api_key=GEMINI_API_KEY)
构建电影向量索引
本节介绍如何准备电影数据集、使用 Gemini 生成嵌入向量，并将其索引到 Qdrant 中以进行相似性搜索。
1. 从 Kaggle 加载数据集
首先，使用 kagglehub 库从 Kaggle 加载数据集。本笔记本使用的数据集是 TMDB 电影数据集 2024，其中包含约 100 万+ 条电影记录。
import kagglehub
from kagglehub import KaggleDatasetAdapter
file_path = "TMDB_movie_dataset_v11.csv"
df = kagglehub.load_dataset(
KaggleDatasetAdapter.PANDAS,
"asaniczka/tmdb-movies-dataset-2023-930k-movies",
file_path,
)
2. 检查数据集结构
由于数据集很大，提前检查有助于你识别有用的字段并过滤掉不相关的数据。
print("\nDataset Columns:")
print(df.columns)
print("\nMissing Values per Column:")
print(df.isnull().sum())
print(f"\nNumber of rows: {len(df)}")
print(f"Number of unique IDs: {df['id'].nunique()}")
3. 过滤并清洗数据集
此步骤筛选数据集，仅保留对语义搜索有用的元数据：id、title、overview、genres、keywords、tagline 和 release_date。这些字段提供了足够的上下文来生成有意义的嵌入。
缺少 title 或同时缺少 overview 和 genres 的条目（行）将被移除，因为它们没有足够的描述性数据来进行准确的推荐。
import pandas as pd
import numpy as np
import ast
print(f"Original rows: {len(df)}")
columns_to_keep = ['id', 'title', 'overview', 'genres', 'keywords', 'tagline', 'release_date']
df_relevant = df[columns_to_keep].copy()
print(f"Rows before dropping missing title: {len(df_relevant)}")
df_relevant.dropna(subset=['title'], inplace=True)
df_relevant = df_relevant[~(df_relevant['genres'].isna() & df_relevant['overview'].isna())]
print(f"Rows after dropping missing title and dropping missing (genres and overview): {len(df_relevant)}")
# Fill missing text columns with empty strings
text_cols_to_fill = ['overview', 'genres', 'keywords', 'tagline']
for col in text_cols_to_fill:
df_relevant[col] = df_relevant[col].fillna('')
# Extract release year from the release_date string
def get_year(date_str):
if pd.isna(date_str) or not isinstance(date_str, str) or len(date_str) < 4:
return None
try:
return int(date_str[:4])
except (ValueError, TypeError):
return None
df_relevant['release_year'] = df_relevant['release_date'].apply(get_year)
print("\nSample data after cleaning (keeping missing overviews):")
print(df_relevant[['id', 'title', 'overview', 'genres', 'keywords', 'tagline', 'release_year']].head())
4. 准备用于嵌入的文本
此步骤通过将相关字段合并为单个结构化文本字符串，为嵌入准备电影元数据。该表示形式包括标题、概述、类型、关键词、标语和发行年份（如果可用）。输出存储在一个名为 text_for_embedding 的新列中。
嵌入是文本的数值向量表示，可捕获语义含义和关系。这些向量可用于相似性搜索和聚类等任务。
了解更多关于文本嵌入的信息，并探索Gemini 嵌入笔记本。
def create_embedding_text(row):
"""Combines available movie metadata into a single string for embedding."""
# Title is always present, so it can be included directly
title_str = f"Title: {row['title']}"
overview_str = f"Overview: {row['overview']}" if row['overview'] else ""
year_str = f"Release Year: {int(row['release_year'])}" if pd.notna(row['release_year']) else ""
genre_str = f"Genres: {row['genres']}" if row['genres'] else ""
keywords_str = f"Keywords: {row['keywords']}" if row['keywords'] else ""
tagline_str = f"Tagline: {row['tagline']}" if row['tagline'] else ""
parts = [
title_str,
overview_str,
year_str,
genre_str,
keywords_str,
tagline_str
]
return "\n".join(part for part in parts if part)
df_relevant['text_for_embedding'] = df_relevant.apply(create_embedding_text, axis=1)
# Use this to inspect how movie data has been transformed for embedding
print(df_relevant[['id', 'title', 'text_for_embedding']].head())
5. 为开发采样一个子集
为了让笔记本易于运行并确保高效开发，你需要快速迭代并尽量减少资源使用。此步骤不是使用完整数据集，而是从清洗后的数据中抽样5000部电影，除非数据集本身已经更小，在这种情况下则使用所有条目。
SAMPLE_SIZE = 5000
if len(df_relevant) > SAMPLE_SIZE:
print(f"\nTaking a random sample of {SAMPLE_SIZE} movies for development.")
df_sample = df_relevant.sample(n=SAMPLE_SIZE, random_state=42)
else:
print(f"\nCleaned dataset size ({len(df_relevant)}) is smaller than or equal to SAMPLE_SIZE. Using the full cleaned dataset.")
df_sample = df_relevant
print(f"Working with {len(df_sample)} movies for the next steps.")
print(df_sample[['id', 'title', 'release_year']].head())
columns_for_payload = ['title', 'overview', 'genres', 'keywords', 'tagline', 'release_year']
columns_final = ['id', 'text_for_embedding'] + columns_for_payload
df_sample = df_sample[columns_final]
print("\nFinal sample DataFrame structure for embedding/indexing:")
print(df_sample.info())
6. 初始化 Qdrant 以进行向量索引
数据准备就绪后，下一步是设置 Qdrant，这是一个向量相似性搜索引擎，专为存储和查询高维向量而优化。它支持快速索引、过滤以及跨数百万向量的相似性搜索。
Qdrant 可以运行：
* 本地作为独立服务
* 云端用于生产部署
* 或者完全 内存模式，用于开发期间的快速临时使用
在本笔记本中，通过向客户端传递 ":memory:" 来初始化 Qdrant 的内存模式。这会将数据仅存储在 RAM 中，意味着 会话结束后数据不会持久化。这适用于实验，但不适合长期保存结果。
您还需要配置以下内容：
* COLLECTION_NAME：用于存储电影向量的 Qdrant 集合名称
* VECTOR_SIZE：设置为 3072，以匹配 Gemini 生成的文本嵌入的维度
* DISTANCE_METRIC：设置为 余弦距离，这非常适合衡量嵌入向量之间的语义相似性
from qdrant_client import QdrantClient, models
import time
COLLECTION_NAME = "tmdb_movies_sample"
VECTOR_SIZE = 3072
DISTANCE_METRIC = models.Distance.COSINE
# Initialize Qdrant client using in-memory storage
qdrant_client = QdrantClient(":memory:")
7. 定义批量嵌入函数
此步骤定义了 get_embeddings_batch 函数，该函数使用 Gemini 嵌入模型（gemini-embedding-001）为批量电影数据生成文本嵌入，并包含自动重试机制以增强鲁棒性。
import time
from google.api_core import exceptions, retry
MODEL_FOR_EMBEDDING = "gemini-embedding-001" # @param ["gemini-embedding-001", "gemini-embedding-2-preview"] {"allow-input":true, isTemplate: true}
BATCH_SIZE = 25
QDRANT_BATCH_SIZE = 3072
@retry.Retry(timeout=3000)
def get_embeddings_batch(texts: list[str], task_type="RETRIEVAL_DOCUMENT") -> list[list[float]] | None:
"""
Generates embeddings for a batch of texts using Gemini API with retry.
Args:
texts: A list of strings to embed.
task_type: The task type for the embedding model.
Returns:
A list of embedding vectors (list of floats), or None if a non-retryable error occurs.
"""
if not texts:
return []
try:
response = client.models.embed_content(
model=MODEL_FOR_EMBEDDING,
contents=texts,
config={
"task_type":task_type,
}
)
return response.embeddings
except exceptions.RetryError as e:
print(f"Embedding batch failed after retries: {e}")
return None
except Exception as e:
print(f"An unexpected error occurred during embedding: {e}")
return None
# Example of what an embedding looks like
sample_embedding = get_embeddings_batch(["Example movie about space and survival"])[0]
print("Example embedding vector:", sample_embedding.values[:10])
8. 在 Qdrant 中创建集合
在 Qdrant 中，集合（collection）类似于数据库中的表，它存储向量以及可选的元数据（payload）。每个集合都有自己的配置，包括向量大小和相似度度量。
# In case someone tries running the whole notebook again they would want to create the collection again
try:
qdrant_client.delete_collection(collection_name=COLLECTION_NAME)
print(f"Existing collection '{COLLECTION_NAME}' deleted.")
except Exception as e:
print(f"Error deleting collection (it might not exist): {e}")
try:
qdrant_client.create_collection(
collection_name=COLLECTION_NAME,
vectors_config=models.VectorParams(
size=VECTOR_SIZE,
distance=DISTANCE_METRIC
)
)
print(f"Collection '{COLLECTION_NAME}' created successfully.")
except Exception as e:
print(f"Error creating collection: {e}")
9. 创建用于元数据存储的 Payload
在 Qdrant 中，除了存储向量嵌入之外，你还可以将附加信息（称为 payload）附加到每个向量上。这些元数据有助于根据标题、类型或发行年份等属性进行过滤或检索相关结果。
create_payload 函数通过从每条电影记录中提取指定列来准备 payload，处理缺失值并确保数据类型与 Qdrant 兼容。
payload_columns = [
'title', 'overview', 'genres', 'keywords', 'tagline', 'release_year'
]
def create_payload(row, payload_columns):
payload = {}
for col in payload_columns:
value = row[col]
if pd.isna(value):
payload[col] = None
elif isinstance(value, (np.int64, np.int32)):
payload[col] = int(value)
elif isinstance(value, (np.float64, np.float32)):
payload[col] = float(value)
else:
payload[col] = value
return payload
10. 批量嵌入并索引到 Qdrant
此步骤分批处理采样的电影数据集，使用 Gemini API 生成向量嵌入，并将这些嵌入及其元数据负载上传（upsert）到 Qdrant 集合。
此过程的关键点：
* 数据集按 BATCH_SIZE 大小分批处理，以生成嵌入，确保在 API 限制范围内。
* 每批文本数据发送到 Gemini 嵌入 API，嵌入函数中处理重试。
* 对于每个成功嵌入的批次，代码准备 点（每个点包含 ID、向量嵌入和元数据负载）以上传到 Qdrant。
* 点被缓冲并按 QDRANT_BATCH_SIZE 大小分块上传，以优化性能。
* 该过程包含错误处理和重试逻辑，以避免失败导致整个操作停止。
* 最后，缓冲区中剩余的点会被上传。
* 打印已处理、失败和成功 upsert 的项目的汇总统计信息。
from tqdm import tqdm
print(f"Starting batch embedding and indexing process for {len(df_sample)} movies...")
print(f"Using Gemini Batch Size: {BATCH_SIZE}, Qdrant Upsert Batch Size: {QDRANT_BATCH_SIZE}")
points_to_upsert_buffer = []
total_processed = 0
total_failed_embedding = 0
total_upserted = 0
num_batches = (len(df_sample) + BATCH_SIZE - 1) // BATCH_SIZE
for i in tqdm(range(0, len(df_sample), BATCH_SIZE), total=num_batches, desc="Processing Batches"):
batch_df = df_sample.iloc[i : i + BATCH_SIZE]
batch_texts = batch_df['text_for_embedding'].tolist()
batch_ids = batch_df['id'].tolist()
if not batch_texts:
continue
# Generate embeddings for the current batch of movie texts
batch_embeddings = get_embeddings_batch(batch_texts, task_type="RETRIEVAL_DOCUMENT")
# Check if embeddings were successfully generated and correspond to batch size
if batch_embeddings and len(batch_embeddings) == len(batch_texts):
for j in range(len(batch_ids)):
item_id = batch_ids[j]
item_embedding = batch_embeddings[j]
row_data = batch_df.iloc[j]
# Prepare metadata payload for this movie
payload = create_payload(row_data, payload_columns)
# Create a Qdrant PointStruct with id, embedding vector, and payload
point = models.PointStruct(
id=int(item_id),
vector=item_embedding.values,
payload=payload
)
points_to_upsert_buffer.append(point)
total_processed += len(batch_ids)
else:
print(f"Failed to get embeddings for batch starting at index {i}. Skipping {len(batch_ids)} items.")
total_failed_embedding += len(batch_ids)
continue
# Upload buffered points to Qdrant if buffer reached batch size or end of data
if len(points_to_upsert_buffer) >= QDRANT_BATCH_SIZE or (i + BATCH_SIZE >= len(df_sample)):
if points_to_upsert_buffer:
try:
qdrant_client.upsert(
collection_name=COLLECTION_NAME,
points=points_to_upsert_buffer,
wait=False
)
total_upserted += len(points_to_upsert_buffer)
points_to_upsert_buffer = []
except Exception as e:
print(f"Error upserting chunk to Qdrant: {e}")
points_to_upsert_buffer = []
time.sleep(5)
# Pause before retrying to avoid hammering the service after an error
# Upload any remaining points left in buffer after loop completion
if points_to_upsert_buffer:
print(f"Upserting final remaining chunk of {len(points_to_upsert_buffer)} points.")
try:
qdrant_client.upsert(
collection_name=COLLECTION_NAME,
points=points_to_upsert_buffer,
wait=True
)
total_upserted += len(points_to_upsert_buffer)
points_to_upsert_buffer = []
except Exception as e:
print(f"Error upserting final chunk: {e}")
print("Batch embedding and indexing finished.")
print(f"Total items processed (attempted embedding): {total_processed}")
print(f"Total points successfully prepared for upsert: {total_upserted}")
# Waiting for collection to settle
time.sleep(5)
try:
count = qdrant_client.count(collection_name=COLLECTION_NAME, exact=True)
print(f"\nVerification: Collection '{COLLECTION_NAME}' now contains {count.count} points.") # it should print 5000
except Exception as e:
print(f"Error verifying collection count: {e}")
11. 使用向量嵌入搜索和推荐相似电影
将所有电影向量索引到 Qdrant 后，您现在可以执行语义搜索。这使您能够获取任何用户查询（例如短语、电影描述或概念），使用相同的 Gemini 模型将其转换为嵌入，并使用余弦相似度从集合中检索最相似的电影向量。
这个 recommend_movies 函数演示了如何：
* 使用 Gemini API 从您的输入查询生成嵌入。
* 使用 Qdrant 的 search() 方法执行相似性搜索。
* 检索最相似的 k 个电影条目，包括它们的元数据和相似度分数。
这是向量数据库作为推荐引擎的最后一步。
def recommend_movies(query_text, top_k=5):
"""
Finds movies similar to the query_text using the Qdrant index.
Args:
query_text (str): The user's query (e.g., movie title, description, theme).
top_k (int): The maximum number of recommendations to return.
Returns:
list: A list of dictionaries, where each dictionary contains the
payload (movie details) and similarity score of a recommended movie.
Returns an empty list if query embedding fails or no results found.
"""
print(f"Searching for recommendations based on: '{query_text}'")
# Generate embedding for the user query using Gemini
query_embedding = get_embeddings_batch(query_text, task_type="RETRIEVAL_QUERY")[0].values
print(f"Query embedding: {query_embedding}")
if query_embedding is None:
print("Error: Could not generate embedding for the query.")
return []
try:
# Perform a semantic search on Qdrant using the query vector
search_result = qdrant_client.search(
collection_name=COLLECTION_NAME,
query_vector=query_embedding,
limit=top_k,
with_payload=True
)
recommendations = []
if search_result:
print(f"Found {len(search_result)} potential recommendations:")
for hit in search_result:
recommendation = {
"id": hit.id,
"score": hit.score,
"payload": hit.payload
}
recommendations.append(recommendation)
else:
print("No recommendations found matching the query.")
return recommendations
except Exception as e:
print(f"Error during Qdrant search: {e}")
return []
试试你的电影推荐器
现在你可以通过自然语言描述一个主题、类型或概念来查询你的电影推荐器。系统将基于使用 Gemini 生成的嵌入向量进行向量相似性搜索，从你的数据集中返回语义上最相似的电影。
query = """
I want to watch something with my girlfriends that’s
both funny and teaches something.
"""
recommendations = recommend_movies(query, top_k=5)
if recommendations:
print("\n--- Recommendations ---")
for rec in recommendations:
print(f" - Score: {rec['score']:.4f}")
print(f" Title: {rec['payload'].get('title', 'N/A')}")
print(f" Genre: {rec['payload'].get('genres', 'N/A')}")
print(f" Year: {rec['payload'].get('release_year', 'N/A')}")
print("-" * 10)
后续步骤
本笔记本演示了如何通过结合 Gemini API 的嵌入功能与 Qdrant 的向量搜索来构建电影推荐系统。
有用的 API 参考
如需更详细的理解并探索高级功能，请参阅官方文档：
* Gemini API 嵌入文档：了解如何生成文本嵌入、理解模型参数，并将其有效用于语义搜索和相似性任务。
* Qdrant Python 客户端文档：了解如何管理集合、插入和搜索向量、配置索引以及与向量数据库交互。
相关示例
要探索更多用例并获得额外灵感，请查看此目录中的相关笔记本：
* 使用 Qdrant 进行相似性搜索：一个专注于使用嵌入和 Qdrant 构建语义搜索系统的示例。
* Google GenAI SDK 概述：引导您完成 SDK 的安装和设置、文本和多模态提示、令牌计数、安全过滤器、多轮聊天、函数调用、文件上传、上下文缓存等。
* 使用 Gemini API 进行文本嵌入：专注于使用 Gemini API 生成和处理文本嵌入，非常适合构建基于向量的搜索和推荐系统。
常见问题
常见错误：1) API密钥未设置，导致认证失败——检查GEMINI_API_KEY环境变量。2) 嵌入维度不匹配——确保VECTOR_SIZE与模型输出一致（3072）。3) Qdrant集合已存在——先删除再创建。4) 批量嵌入失败——检查网络或配额，重试机制已内置。5) 内存不足——减少SAMPLE_SIZE或使用磁盘存储。排查方法：查看日志输出，验证集合计数，测试单个嵌入。
企业应用扩展
优化方向：1) 使用Qdrant云服务实现持久化与扩展。2) 引入混合检索（结合BM25）提升精度。3) 增加用户反馈机制，动态调整推荐。4) 使用更高效的嵌入模型（如gemini-embedding-2-preview）。5) 部署为REST API，集成到Web应用。6) 利用Qdrant的payload过滤功能，按类型、年份等条件筛选。
总结
本教程面向有编程基础的开发者，深入讲解如何利用Gemini API生成文本嵌入，并结合Qdrant向量数据库构建电影推荐系统。内容涵盖环境配置、数据清洗、批量嵌入、向量索引及语义搜索，并提供完整代码与架构分析。通过本教程，读者将掌握构建生产级推荐系统的核心技能，包括处理大规模数据、优化API调用及调试常见问题。
- --
内容来源与验证
- 内容等级：PREMIUM_A
- 质量等级：S_CANDIDATE
- 难度：中级
- 验证等级：V1
- 原文作者：
- 原文日期：
- 来源许可：Apache-2.0
- 来源版本：
addad0f04a2e779b9e50dd344b4a337fb359fcf6 - 查看原文
站内还有更多完整实战教程
加入会员，解锁全部教程、完整代码与持续更新内容。
加入会员