一文简介 Bert 模型、向量表示、向量搜索、向量库、API 构建、API 接口、并行计算、RAG、Next.js、React.js、LangChain

We recommend cosine similarity. The choice of distance function typically doesn't matter much.

OpenAI embeddings are normalized to length 1, which means that:
    - Cosine similarity can be computed slightly faster using just a dot product
    - Cosine similarity and Euclidean distance will result in the identical rankings

import pandas as pd 
import numpy as np 
import uuid

df = pd.read_csv('../static/files/tmdb-movies.csv')
df.shape

df.columns 

df[['id', 'title', 'overview', 'revenue', 'poster_path']].head(5)

movie_ids = df.id.tolist()
movie_overviews = df.overview.tolist()
# Create dictionaries
Id2Title = dict(zip(df.id, df.original_title))
Id2Overview = dict(zip(df.id, df.overview))
Id2PosterPath = dict(zip(df.id, df.poster_path))
Id2Revenue = dict(zip(df.id, df.revenue))

from sentence_transformers import SentenceTransformer
from sklearn.preprocessing import normalize
model = SentenceTransformer('all-MiniLM-L6-v2')
from typing import List, Dict, Tuple
import faiss 

def get_movie_embeddings(
    movie_ids:List[str], 
    movie_overviews:List[str], 
    model: SentenceTransformer,
    batch_size: int = 32, 
    ) -> Dict[str, np.ndarray]:
    """
    Generates normalized embeddings for movie overviews.

    Args:
        movie_ids: List of movie IDs.
        movie_overviews: List of movie overviews.
        model: Embedding model (e.g., SentenceTransformer).
        batch_size: Batch size for processing.

    Returns:
        Dictionary mapping movie IDs to normalized embeddings.
    """
    if len(movie_ids) != len(movie_overviews):
        raise ValueError("movie_ids and movie_overviews must have the same length.")

    all_embeddings = []
    for i in range(0, len(movie_ids), batch_size):
        batch_movies = movie_overviews[i:i+batch_size]
        try:
            batch_embeddings = model.encode(batch_movies)
            all_embeddings.extend(batch_embeddings)
        except Exception as e:
            print(f"Error encoding batch {i}: {e}")
            raise
    try:
        all_embeddings_np = np.array(all_embeddings) #added numpy array conversion.
        normalized_embeddings = normalize(all_embeddings_np, axis=1, norm='l2')
    except Exception as e:
        print(f"Error normalizing embeddings: {e}")
        raise
    # Create a hashmap; otherwise we cannot map from embeddings to distinct movies
    # starting from Python 3.7, dictionaries (dict) preserve insertion order by default.
    # Therefore, there is no need to sort
    movie_embeddings = dict(zip(movie_ids, normalized_embeddings))
    return movie_embeddings

movie_embeddings = get_movie_embeddings(movie_ids, movie_overviews, model)

def prepare_faiss_index(
    movie_embeddings: Dict[str, np.ndarray]) -> Tuple[faiss.Index, List[str]]:
    """
    Prepares a FAISS index from movie embeddings.
    
    Args:
        movie_embeddings: Dictionary of movie IDs to normalized embeddings.
    
    Returns:
        Tuple of (FAISS index, ordered list of movie IDs)
    """
    movie_ids = list(movie_embeddings.keys())
    embeddings = np.array([movie_embeddings[mid] for mid in movie_ids])

    #create and populate the index
    dim = embeddings.shape[1]
    index = faiss.IndexFlatIP(dim)
    index.add(embeddings)

    return index, movie_ids

def search_movies(
    faiss_index: faiss.Index,
    movie_ids: List[str],
    id2title_dict: Dict[int, str],
    user_query: str,
    model: SentenceTransformer,
    k: int = 10,
) -> Tuple[List[str], List[float]]:
    """
    Searches movies based on a user query using Faiss and cosine similarity.

    Args:
        faiss_index: Pre-built FAISS index
        movie_ids: Ordered list of movie IDs that corresponds to the index
        id2title_dict: Dict mapping from movie IDs to movie titles
        user_query: The user's search query.
        model: SentenceTransformer model.
        k: Number of results to return.

    Returns:
        Tuple, top K movie titles and their associated cosine similarity scores
    """
    #embed and normalize the user query
    query_embedding = model.encode([user_query])
    normalized_query_embedding = normalize(query_embedding, axis = 1, norm = 'l2')

    #search the index
    similarities, indices = faiss_index.search(normalized_query_embedding, k)

    #retrieve movie ids and similarity scores
    top_k_ids = [movie_ids[i] for i in indices[0]]
    top_k_titles = [id2title_dict[x] for x in top_k_ids]
    top_k_scores = similarities[0]
    return top_k_titles, top_k_scores

faiss_index, movie_ids = prepare_faiss_index(movie_embeddings)
user_query = "NASA tried to rescue an astronaut stranded on Mars."
results = search_movies(
    faiss_index,
    movie_ids,
    Id2Title,
    user_query,
    model
)

results 

pip install qdrant-client

import os 
# `pip install python-dotenv` first
from dotenv import load_dotenv
load_dotenv()
url = os.getenv("QDRANT_MOVIE")
api_key = os.getenv("QDRANT_API_KEY")

from qdrant_client import QdrantClient
from qdrant_client.models import VectorParams, Distance
import uuid

qdrant_client = QdrantClient(
    url=url,
    api_key=api_key
)

try:
    qdrant_client.delete_collection(collection_name='movies')
    print("Deleted existing collection")
except Exception as e:
    print(f"Collection might not exist yet: {e}")

# Create the collection first
vector_size = len(next(iter(movie_embeddings.values())))
qdrant_client.create_collection(
    collection_name="movies",
    vectors_config=VectorParams(
        size=vector_size,
        distance=Distance.COSINE
    )
)

# Qdrant requires ids to be strs, not int
def int2uuid(i:int):
    return str(uuid.uuid5(uuid.NAMESPACE_DNS, str(i)))

batch_size = 1_000 
for i in range(0, len(movie_ids), batch_size):
    batch_ids = movie_ids[i:i+batch_size]
    batch_vectors = []
    for movie_id in batch_ids:
        batch_vectors.append({
        'id': movie_id,
        'vector': movie_embeddings[movie_id].tolist(),
        'payload': {
            'title': Id2Title.get(movie_id, ''),
            'overview': Id2Overview.get(movie_id, ""),
            'poster_path': Id2PosterPath.get(movie_id, ""),
            'revenue': Id2Revenue.get(movie_id, "")
        }
    })
    qdrant_client.upsert(
        collection_name = 'movies',
        points = batch_vectors
    )
    print(f"Uploaded batch {i//batch_size + 1}/{(len(movie_ids)-1)//batch_size + 1}")

def search_qdrant(
    query_text:str, model:SentenceTransformer, top_k:int = 5):
    query_embedding = model.encode([query_text])
    normalized_query = normalize(query_embedding, axis = 1, norm='l2')

    results = qdrant_client.query_points(
        collection_name = 'movies',
        query=normalized_query[0].tolist(),
        limit=top_k
    )
    movies = []
    for result in results.points:
        movies.append({
            'title': result.payload['title'],
            'overview': result.payload['overview'],
            'similarity_score': result.score
        })
    return movies 

user_query = "NASA tried to rescue an astronaut stranded on Mars."
search_qdrant(query_text=user_query, model=model)

pip install pinecone

from pinecone import Pinecone, ServerlessSpec

pinecone_api_key = api_key = os.getenv("PINECONE_API_KEY")
pc = Pinecone(api_key=pinecone_api_key)

index_name = "movies"

# Get the list of index names from the dictionaries
existing_indexes = [index["name"] for index in pc.list_indexes()]

if index_name in existing_indexes:
    print(f"Deleting existing index: {index_name}!")
    pc.delete_index(index_name)
else:
    print(f"Index '{index_name}' does not exist.")

vector_size = len(next(iter(movie_embeddings.values())))

pc.create_index(
    name=index_name,
    dimension=vector_size, # Replace with your model dimensions
    metric="cosine", # Replace with your model metric
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    ) 
)

# Get the index object
index = pc.Index(index_name)

batch_size = 1_000
for i in range(0, len(movie_ids), batch_size):
    batch_ids = movie_ids[i:i+batch_size]
    batch_vectors = []
    for movie_id in batch_ids:
        batch_vectors.append({
            'id': str(movie_id),
            'values': movie_embeddings[movie_id].tolist(),
            'metadata':{
                'title': Id2Title.get(movie_id, ''),
                'overview': Id2Overview.get(movie_id, ""),
                'poster_path': Id2PosterPath.get(movie_id, ""),
                'revenue': Id2Revenue.get(movie_id, "")
            }
        })
    index.upsert(vectors = batch_vectors)
    print(f"Uploaded batch {i//batch_size + 1}/{(len(movie_ids)-1)//batch_size + 1}")

def search_pinecone(
    query_text:str, 
    model:SentenceTransformer, 
    top_k:int = 5):
    query_embedding = model.encode([query_text])
    normalized_query = normalize(query_embedding, axis = 1, norm='l2')

    results = index.query(
        vector=normalized_query[0].tolist(),
        top_k=top_k,
        include_metadata=True
    )

    # Format results
    movies = []
    for match in results['matches']:
        movies.append({
            "title": match['metadata']['title'],
            "overview": match['metadata']['overview'],
            "similarity_score": match['score']
        })
    
    return movies

user_query = "NASA tried to rescue an astronaut stranded on Mars."
search_pinecone(query_text=user_query, model=model)

fastapi[standard]
sentence_transformers
qdrant_client
python-dotenv==1.0.0
scikit-learn==1.3.2

import google.generativeai as genai
import os

GEMINI_API_KEY = os.getenv("GEMINI_API_KEY")

genai.configure(api_key=GEMINI_API_KEY)

norm = np.linalg.norm(embedding_model['embedding'])

print(f"L2 范数：{norm}")
if np.isclose(norm, 1):
    print("向量已进行 L2 归一化")
else:
    print("向量未进行 L2 归一化")

def get_gemini_embeddings(texts:List[str]) -> np.ndarray:
    embeddings = []
    for text in texts:
        response = genai.embed_content(
            model="models/text-embedding-004", 
            content=text)
        embeddings.append(response['embedding'])
    return np.array(embeddings)

try:
    qdrant_client.delete_collection(collection_name='movies_genai')
    print("Deleted existing collection")
except Exception as e:
    print(f"Collection might not exist yet: {e}")

# Create the collection first
embeddings = get_gemini_embeddings([movie_overviews[0]])
vector_size = len(embeddings[0])
qdrant_client.create_collection(
    collection_name="movies_genai",
    vectors_config=VectorParams(
        size=vector_size,
        distance=Distance.COSINE
    )
)

from concurrent.futures import ThreadPoolExecutor
import threading 
from tqdm import tqdm 

def process_one_batch(batch_ids, batch_overviews):
    try:
        batch_embeddings = get_gemini_embeddings(batch_overviews)
        points = []
        for j, movie_id in enumerate(batch_ids):
            points.append({
                'id': movie_id,
                'vector': batch_embeddings[j].tolist(),
                'payload': {
                    'title': Id2Title.get(movie_id, ''),
                    'overview': Id2Overview.get(movie_id, ""),
                    'poster_path': Id2PosterPath.get(movie_id, ""),
                    'revenue': Id2Revenue.get(movie_id, "")
                }
            })
        qdrant_client.upsert(
            collection_name = 'movies_genai',
            points = points
        )
        return True 
    except Exception as e:
        print(f"Error uploading batch: {e}")
        return False 

# #|code-fold:true

# batch_size = 180
# all_batches = []
# for i in range(0, len(movie_ids), batch_size):
#     batch_ids = movie_ids[i:i+batch_size]
#     batch_overviews = movie_overviews[i:i + batch_size]
#     all_batches.append((batch_ids, batch_overviews))

# max_workers = 3 # 同时进行 3 个线程
# with ThreadPoolExecutor(max_workers=max_workers) as executor:
#     results = list(tqdm(
#         executor.map(lambda batch: process_one_batch(batch[0], batch[1]), all_batches),
#         total=len(all_batches),
#         desc="Uploading batches"
#     ))

# # 统计成功数量
# successful_batches = results.count(True)
# print(f"成功处理 {successful_batches}/{len(all_batches)} 批次")

def search_qdrant(
    query_text:str, top_k:int = 5):
    # 使用 Google-Gemini 而不是我们之前的 model
    query_embeddings = get_gemini_embeddings([query_text])

    results = qdrant_client.query_points(
        collection_name = 'movies_genai',
        query=query_embeddings[0].tolist(),
        limit=top_k
    )
    movies = []
    for result in results.points:
        movies.append({
            'title': result.payload['title'],
            'overview': result.payload['overview'],
            'similarity_score': result.score
        })
    return movies 

user_query = "NASA tried to rescue an astronaut stranded on Mars."
search_qdrant(query_text=user_query)

user_query = "NASA尝试营救一名滞留在火星上的宇航员。"
search_qdrant(query_text=user_query)

df[['id', 'overview', 'revenue', 'title', 'poster_path']].to_json(
        './data/movies.json', orient='records', indent=4)

node -v
npm -v

npx create-next-app@latest my-next-app

cd my-next-app