使用 MMR 搜索平衡相关性和多样性

pip install qdrant-client # Vector Search Engine
pip install fastembed # Fast, lightweight embedding generation with CLIP
pip install datasets # For DeepFashion data access

from datasets import load_dataset
from fastembed import ImageEmbedding, TextEmbedding
from qdrant_client import QdrantClient, models
import uuid

def load_fashion_data(sample_size=10):
    """Load fashion items from DeepFashion-with-masks dataset"""
    dataset = load_dataset("SaffalPoosh/deepFashion-with-masks", split="train")
    sample_dataset = dataset.shuffle(seed=42).select(range(sample_size))
    
    fashion_items = []
    for i, item in enumerate(sample_dataset):
        # Extract all available metadata
        metadata = {}
        for key, value in item.items():
            if key not in ["images", "mask", "mask_overlay"] and value is not None:
                metadata[key] = value
        
        fashion_items.append({
            "image": item["images"],  # PIL Image object
            **metadata
        })
    
    return fashion_items

fashion_items = load_fashion_data(sample_size=100)

# Create image embeddings for all fashion items
image_model = ImageEmbedding(model_name="Qdrant/clip-ViT-B-32-vision")
embeddings = list(image_model.embed([item["image"] for item in fashion_items]))

# Initialize Qdrant client, get your credentials at https://qdrant.org.cn
client = QdrantClient(
     host="your-qdrant url",
     api_key="<your-qdrant-api-key>")

collection_name = "fashion_discovery"

# Create collection optimized for CLIP embeddings
client.recreate_collection(
    collection_name=collection_name,
    vectors_config=models.VectorParams(
        size=512,  # CLIP embedding dimension
        distance=models.Distance.COSINE
    )
)

# Upload fashion items with embeddings
points = []
for i, (item, embedding) in enumerate(zip(fashion_items, embeddings)):
    # Remove PIL image from payload
    payload = {k: v for k, v in item.items() if k != "image"}
    payload["item_id"] = i
    
    point = models.PointStruct(
        id=str(uuid.uuid4()),
        vector=embedding.tolist(),
        payload=payload
    )
    points.append(point)

client.upsert(collection_name=collection_name, points=points)

# Standard fashion search — often returns similar looking items

def fashion_search_standard(query_text, limit=5):
    text_model = TextEmbedding(model_name="Qdrant/clip-ViT-B-32-text")
    query_embedding = list(text_model.embed([query_text]))[0]
    
    results = client.query_points(
        collection_name=collection_name,
        query=query_embedding.tolist(),
        limit=limit,
        with_payload=True
    )
    return results

query_text = "black jacket"

# Standard search
standard_results = fashion_search_standard(query_text)
print("\nSTANDARD FASHION SEARCH: 'black jacket'")
for i, point in enumerate(standard_results, 1):
    payload = point.payload
    print(f"{i}️⃣ Score: {point.score:.4f} | {payload.get('item_description', 'Fashion Item')}")
    print(f"   Style: {payload.get('style', 'N/A')} | Color: {payload.get('color', 'N/A')}")
    print()

 # MMR fashion search — balances relevance with style diversity

def fashion_search_mmr(query_text, limit=5, diversity=0.5):
    text_model = TextEmbedding(model_name="Qdrant/clip-ViT-B-32-text")
    query_embedding = list(text_model.embed([query_text]))[0]
    
    results = client.query_points(
        collection_name=collection_name,
        query=models.NearestQuery(
            nearest=query_embedding.tolist(),
            mmr=models.Mmr(
                diversity=diversity,  # 0.0 - relevance; 1.0 - diversity
                candidates_limit=100  # num of candidates to preselect
            )
        ),
        limit=limit,
        with_payload=True
    )
    return results

query_text = "black jacket"

# MMR search with diversity=0.5
mmr_results = fashion_search_mmr(query_text, diversity=0.5)
print("\nMMR FASHION SEARCH: 'black jacket' (diversity=0.5)")
for i, point in enumerate(mmr_results.points, 1):
    payload = point.payload
    print(f"{i}️⃣ Score: {point.score:.4f} | {payload.get('item_description', 'Fashion Item')}")
    print(f"   Style: {payload.get('style', 'N/A')} | Color: {payload.get('color', 'N/A')}")
    print()

def filtered_fashion_search(query_text, metadata_filter=None, limit=5, diversity=0.4):
    text_model = TextEmbedding(model_name="Qdrant/clip-ViT-B-32-text")
    query_embedding = list(text_model.embed([query_text]))[0]

    # Build filter for metadata fields
    filter_conditions = []
    if metadata_filter:
        for key, value in metadata_filter.items():
            filter_conditions.append(
                models.FieldCondition(
                    key=key,
                    match=models.MatchValue(value=value)
                )
            )

    query_filter = models.Filter(must=filter_conditions) if filter_conditions else None

    results = client.query_points(
        collection_name=collection_name,
        query=models.NearestQuery(
            nearest=query_embedding.tolist(),
            mmr=models.Mmr(
                diversity=diversity,  # 0.0 - relevance; 1.0 - diversity
                candidates_limit=100  # num of candidates to preselect
            )
        ),
        query_filter=query_filter,
        limit=limit,
        with_payload=True
    )
    return results

# Example: Filter for women's blouses 
results = filtered_fashion_search(
    "professional attire",
    metadata_filter={"gender": "WOMEN", "cloth_type": "Blouses_Shirts"},
    diversity=0.5)

print("FILTERED SEARCH (Women's Professional Attire):")
for point in results.points:
    payload = point.payload
    print(f"Score: {point.score:.4f}")
    print(f"Gender: {payload.get('gender', 'N/A')}")
    print(f"Cloth Type: {payload.get('cloth_type', 'N/A')}")
    print()