多模态检索与推荐

import numpy as np
from numpy.linalg import norm

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """计算两个向量的余弦相似度
    返回值在 [-1, 1] 之间，1 表示完全相同方向
    """
    dot = np.dot(a, b)
    norm_a = norm(a)
    norm_b = norm(b)
    if norm_a == 0 or norm_b == 0:
        return 0.0
    return float(dot / (norm_a * norm_b))

def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """计算欧氏距离
    对归一化向量，欧氏距离和余弦相似度等价
    d = sqrt(2 * (1 - cos))
    """
    return float(np.sqrt(np.sum((a - b) ** 2)))

# 多模态相似度计算示例
# image_embed: [512] 从 CLIP 图像编码器获得
# text_embed:  [512] 从 CLIP 文本编码器获得
image_embed = np.random.randn(512).astype(np.float32)
text_embed = np.random.randn(512).astype(np.float32)
sim = cosine_similarity(image_embed, text_embed)
print(f"跨模态相似度: {sim:.4f}")

# 暴力检索（Brute-Force）作为基准
import numpy as np
import time

class BruteForceSearcher:
    """暴力检索器：计算查询向量与所有库向量的相似度"""

    def __init__(self, dim: int, metric: str = "cosine"):
        self.dim = dim
        self.metric = metric
        self.vectors = None
        self.ids = None

    def add(self, vectors: np.ndarray, ids: list):
        """添加向量到索引"""
        self.vectors = vectors.astype(np.float32)
        self.ids = ids
        if self.metric == "cosine":
            # 预归一化
            norms = norm(self.vectors, axis=1, keepdims=True)
            norms[norms == 0] = 1.0
            self.vectors = self.vectors / norms

    def search(self, query: np.ndarray, top_k: int = 10):
        """搜索最近邻"""
        query = query.astype(np.float32).reshape(1, -1)
        if self.metric == "cosine":
            query = query / norm(query, axis=1, keepdims=True)
            scores = query @ self.vectors.T  # [1, N]
        else:
            # 欧氏距离
            diff = self.vectors - query  # [N, D]
            scores = -np.sum(diff ** 2, axis=1).reshape(1, -1)

        # Top-K
        indices = np.argpartition(scores[0], -top_k)[-top_k:]
        top_indices = indices[np.argsort(-scores[0, indices])]
        return [(self.ids[i], float(scores[0, i])) for i in top_indices]

# 性能参考: 100 万条 512 维向量，暴力检索约 50-200ms
# 超过 1000 万条时延迟显著增加，需要 ANN 加速

import numpy as np
from sklearn.cluster import KMeans

class ProductQuantizer:
    """乘积量化（PQ）实现"""

    def __init__(self, dim: int, M: int, nbits: int = 8):
        """
        dim: 向量维度
        M: 子空间数量
        nbits: 每个子空间的编码位数
        每个子空间有 2^nbits 个聚类中心
        """
        self.dim = dim
        self.M = M
        self.nbits = nbits
        self.K = 2 ** nbits  # 每个子空间的聚类数
        self.sub_dim = dim // M
        self.codebooks = np.zeros((M, self.K, self.sub_dim))

    def train(self, vectors: np.ndarray):
        """在每个子空间上训练 K-means"""
        N = vectors.shape[0]
        vectors = vectors.reshape(N, self.M, self.sub_dim)
        for m in range(self.M):
            kmeans = KMeans(n_clusters=self.K, n_init=10)
            kmeans.fit(vectors[:, m, :])
            self.codebooks[m] = kmeans.cluster_centers_

    def encode(self, vectors: np.ndarray) -> np.ndarray:
        """将向量编码为 PQ 码本索引"""
        N = vectors.shape[0]
        vectors = vectors.reshape(N, self.M, self.sub_dim)
        codes = np.zeros((N, self.M), dtype=np.uint8)

        for m in range(self.M):
            # 计算到所有聚类中心的距离
            sub_vectors = vectors[:, m, :]  # [N, sub_dim]
            dists = np.sum(
                (sub_vectors[:, np.newaxis, :] - self.codebooks[m][np.newaxis, :, :]) ** 2,
                axis=2
            )  # [N, K]
            codes[:, m] = np.argmin(dists, axis=1)

        return codes

    def compute_distances(self, query: np.ndarray, codes: np.ndarray) -> np.ndarray:
        """非对称距离计算：查询使用原始向量，数据库使用编码"""
        N = codes.shape[0]
        query = query.reshape(self.M, self.sub_dim)
        distances = np.zeros(N)

        for m in range(self.M):
            query_sub = query[m]  # [sub_dim]
            # 查询到该子空间所有聚类中心的距离（查表优化）
            dists = np.sum((self.codebooks[m] - query_sub) ** 2, axis=1)  # [K]
            # 累加对应的码本距离
            distances += dists[codes[:, m]]

        return distances

# HNSW 算法概念模拟
import numpy as np
import heapq

class HNSWNode:
    """HNSW 图中的节点"""

    def __init__(self, vector: np.ndarray, node_id: int, max_layer: int):
        self.vector = vector
        self.node_id = node_id
        # 每层的邻居列表
        self.connections = [[] for _ in range(max_layer + 1)]
        self.max_layer = max_layer

class SimpleHNSW:
    """简化的 HNSW 实现（概念演示）"""

    def __init__(self, dim: int, M: int = 16, ef_construction: int = 200):
        self.dim = dim
        self.M = M
        self.ef_construction = ef_construction
        self.nodes: list[HNSWNode] = []
        self.entry_point = None
        self.max_layer = 0

    def _random_level(self, mL: float = 1.0 / np.log(M)) -> int:
        """根据几何分布随机生成层数"""
        level = 0
        while np.random.random() < 1.0 / mL and level < 16:
            level += 1
        return level

    def _search_layer(self, query: np.ndarray, entry: HNSWNode, ef: int, layer: int):
        """在指定层执行贪心搜索"""
        visited = {entry.node_id}
        candidates = [(-np.dot(entry.vector, query), entry)]  # max-heap
        results = candidates[:]

        while candidates:
            _, current = heapq.heappop(candidates)
            # 最远距离
            worst_dist = -results[0][0] if results else float("inf")
            if np.dot(current.vector, query) < worst_dist:
                break
            for neighbor_id in current.connections[layer]:
                if neighbor_id not in visited:
                    visited.add(neighbor_id)
                    neighbor = self.nodes[neighbor_id]
                    dist = -np.dot(neighbor.vector, query)
                    heapq.heappush(candidates, (dist, neighbor))
                    if len(results) < ef:
                        heapq.heappush(results, (dist, neighbor))
                    elif dist > results[0][0]:
                        heapq.heapreplace(results, (dist, neighbor))

        return results

    def search(self, query: np.ndarray, top_k: int = 10):
        """多层 HNSW 搜索"""
        current = self.entry_point
        # 从上到下逐层搜索
        for layer in range(current.max_layer, 0, -1):
            candidates = self._search_layer(query, current, 1, layer)
            if candidates:
                current = candidates[0][1]
        # 最底层精细搜索
        final_candidates = self._search_layer(query, current, ef=top_k * 2, layer=0)
        final_candidates.sort(key=lambda x: x[0], reverse=True)
        return [(self.nodes[i].node_id, float(d)) for d, node in final_candidates[:top_k]
                for i in [node.node_id]]

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultimodalAligner(nn.Module):
    """多模态嵌入对齐模块
    将不同模态的编码器输出投影到共享空间
    """

    def __init__(self, image_dim: int, text_dim: int, shared_dim: int = 512):
        super().__init__()
        # 图像投影头
        self.image_proj = nn.Sequential(
            nn.Linear(image_dim, shared_dim),
            nn.LayerNorm(shared_dim),
            nn.ReLU(),
            nn.Linear(shared_dim, shared_dim)
        )
        # 文本投影头
        self.text_proj = nn.Sequential(
            nn.Linear(text_dim, shared_dim),
            nn.LayerNorm(shared_dim),
            nn.ReLU(),
            nn.Linear(shared_dim, shared_dim)
        )
        # 可学习的温度参数
        self.logit_scale = nn.Parameter(torch.ones([]) * torch.log(torch.tensor(1.0 / 0.07)))

    def forward(self, image_features: torch.Tensor, text_features: torch.Tensor):
        """计算对比损失"""
        # 投影到共享空间
        image_embeds = F.normalize(self.image_proj(image_features), dim=-1)
        text_embeds = F.normalize(self.text_proj(text_features), dim=-1)

        # 相似度矩阵
        scale = self.logit_scale.exp()
        logits = scale * image_embeds @ text_embeds.T
        batch_size = image_embeds.size(0)
        labels = torch.arange(batch_size, device=image_embeds.device)

        # 对称对比损失
        loss_image = F.cross_entropy(logits, labels)
        loss_text = F.cross_entropy(logits.T, labels)
        return (loss_image + loss_text) / 2

    def get_similarity(self, image_features: torch.Tensor, text_features: torch.Tensor):
        """获取跨模态相似度"""
        image_embeds = F.normalize(self.image_proj(image_features), dim=-1)
        text_embeds = F.normalize(self.text_proj(text_features), dim=-1)
        return image_embeds @ text_embeds.T

# 模态鸿沟量化与分析
import torch
import numpy as np

def measure_modality_gap(image_embeddings: torch.Tensor, text_embeddings: torch.Tensor):
    """量化评估模态鸿沟的多个维度"""

    # 1. 中心距离：两个模态均值向量之间的距离
    image_mean = image_embeddings.mean(dim=0)
    text_mean = text_embeddings.mean(dim=0)
    center_dist = torch.norm(image_mean - text_mean, p=2).item()

    # 2. 方差比率：两个模态的方差是否一致
    image_var = image_embeddings.var(dim=0).mean().item()
    text_var = text_embeddings.var(dim=0).mean().item()
    var_ratio = image_var / (text_var + 1e-8)

    # 3. 类间可分性：同类样本是否跨模态对齐
    # 使用类内相似度 vs 类间相似度
    image_norm = F.normalize(image_embeddings, dim=-1)
    text_norm = F.normalize(text_embeddings, dim=-1)
    cross_sim = torch.diag(image_norm @ text_norm.T).mean().item()

    # 4. 各向异性：嵌入空间的分布均匀性
    def anisotropy(embeddings: torch.Tensor) -> float:
        normalized = F.normalize(embeddings, dim=-1)
        mean_vec = normalized.mean(dim=0)
        return float(torch.norm(mean_vec, p=2))

    image_aniso = anisotropy(image_embeddings)
    text_aniso = anisotropy(text_embeddings)

    return {
        "center_distance": center_dist,
        "variance_ratio": var_ratio,
        "cross_modal_similarity": cross_sim,
        "image_anisotropy": image_aniso,
        "text_anisotropy": text_aniso
    }

# 典型结果 (对齐前 vs 对齐后):
# center_distance:     25.3 -> 2.1
# cross_similarity:     0.12 -> 0.78
# image_anisotropy:     0.95 -> 0.45
# text_anisotropy:      0.88 -> 0.42

import torch
import torch.nn.functional as F

class CrossModalRetrieval:
    """跨模态检索系统 (以文搜图 + 以图搜文)"""

    def __init__(self, image_index, text_index, image_model, text_model, device="cpu"):
        """
        image_index: 图像向量索引 (如 FAISS)
        text_index: 文本向量索引
        image_model: 图像编码器
        text_model: 文本编码器
        """
        self.image_index = image_index
        self.text_index = text_index
        self.image_model = image_model.to(device)
        self.text_model = text_model.to(device)
        self.device = device

    @torch.no_grad()
    def text_to_image(self, query_text: str, top_k: int = 10):
        """以文搜图：用文本查询检索最相关的图像"""
        # 编码查询文本
        text_input = self._tokenize_text(query_text)
        text_feat = self.text_model(text_input)
        text_feat = F.normalize(text_feat, dim=-1).cpu().numpy()

        # 在图像索引中搜索
        scores, indices = self.image_index.search(text_feat, top_k)
        return [(int(indices[0, i]), float(scores[0, i])) for i in range(top_k)]

    @torch.no_grad()
    def image_to_text(self, query_image: torch.Tensor, top_k: int = 10):
        """以图搜文：用图像查询检索最相关的文本"""
        image_feat = self.image_model(query_image.to(self.device))
        image_feat = F.normalize(image_feat, dim=-1).cpu().numpy()

        scores, indices = self.text_index.search(image_feat, top_k)
        return [(int(indices[0, i]), float(scores[0, i])) for i in range(top_k)]

    def evaluate(self, test_images, test_texts, ground_truth, k_values=(1, 5, 10)):
        """评估检索性能: Recall@K"""
        results = {"t2i": {}, "i2t": {}}

        # 以文搜图 Recall
        for q_text, gt_indices in zip(test_texts, ground_truth):
            retrieved = self.text_to_image(q_text, top_k=max(k_values))
            retrieved_ids = [r[0] for r in retrieved]
            for k in k_values:
                recall = len(set(retrieved_ids[:k]) & set(gt_indices)) / len(gt_indices)
                results["t2i"][k] = results["t2i"].get(k, []) + [recall]

        # 取平均
        for k in k_values:
            results["t2i"][k] = sum(results["t2i"][k]) / len(results["t2i"][k])

        return results

# 两阶段检索：召回 + 精排
import torch
from transformers import BlipProcessor, BlipForImageTextRetrieval

class TwoStageRetrieval:
    """两阶段跨模态检索系统"""

    def __init__(self, clip_model, clip_processor, blip_model, blip_processor):
        # 阶段 1: CLIP (快速召回)
        self.clip_model = clip_model
        self.clip_processor = clip_processor

        # 阶段 2: BLIP (精排序)
        self.blip_model = blip_model
        self.blip_processor = blip_processor

    def retrieve(self, query_text: str, image_database: list,
                 recall_k: int = 100, rerank_k: int = 10):
        """
        阶段 1: CLIP 召回 top-K 候选
        阶段 2: BLIP 交叉注意力重排序
        """
        # 阶段 1: CLIP 快速召回
        text_input = self.clip_processor(text=query_text, return_tensors="pt")
        text_feat = self.clip_model.get_text_features(**text_input)
        text_feat = F.normalize(text_feat, dim=-1)

        image_feats = []
        for img in image_database:
            img_input = self.clip_processor(images=img, return_tensors="pt")
            img_feat = self.clip_model.get_image_features(**img_input)
            img_feat = F.normalize(img_feat, dim=-1)
            image_feats.append(img_feat)

        all_image_feats = torch.cat(image_feats, dim=0)
        similarities = (text_feat @ all_image_feats.T)[0]

        # 取 top-K 候选
        top_k_indices = similarities.topk(recall_k).indices.tolist()
        candidates = [(image_database[i], float(similarities[i])) for i in top_k_indices]

        # 阶段 2: BLIP 精排
        reranked = []
        for img, clip_score in candidates[:rerank_k * 2]:
            inputs = self.blip_processor(
                images=img, text=query_text, return_tensors="pt"
            )
            with torch.no_grad():
                output = self.blip_model(**inputs)
                blip_score = output.logits.item()

            # 融合 CLIP 和 BLIP 分数
            fused_score = 0.3 * clip_score + 0.7 * blip_score
            reranked.append((img, fused_score))

        reranked.sort(key=lambda x: x[1], reverse=True)
        return reranked[:rerank_k]

import faiss
import numpy as np

class FAISSIndex:
    """FAISS 向量索引封装"""

    def __init__(self, dim: int, index_type: str = "IVF"):
        self.dim = dim
        self.index_type = index_type
        self.index = None
        self.trained = False

    def build(self, vectors: np.ndarray):
        """构建并训练索引"""
        n = vectors.shape[0]

        if self.index_type == "Flat":
            # 精确检索，无需训练
            self.index = faiss.IndexFlatIP(self.dim)
            self.index.add(vectors)
            self.trained = True

        elif self.index_type == "IVF":
            # IVF 倒排索引
            nlist = min(4096, int(np.sqrt(n)))
            quantizer = faiss.IndexFlatIP(self.dim)
            self.index = faiss.IndexIVFFlat(quantizer, self.dim, nlist, faiss.METRIC_INNER_PRODUCT)

            # 训练
            self.index.train(vectors)
            self.index.add(vectors)
            self.trained = True
            print(f"IVF 索引: {nlist} 个聚类, {n} 个向量")

        elif self.index_type == "IVF_PQ":
            # IVF + 乘积量化（压缩存储）
            nlist = min(4096, int(np.sqrt(n)))
            M = 8  # 子空间数量
            nbits = 8  # 每子空间编码位数
            quantizer = faiss.IndexFlatIP(self.dim)
            self.index = faiss.IndexIVFPQ(quantizer, self.dim, nlist, M, nbits)

            self.index.train(vectors)
            self.index.add(vectors)
            self.trained = True
            print(f"IVF+PQ 索引: {nlist} 聚类, M={M}, 压缩比 {self.dim * 4 / (M + 8):.0f}x")

    def search(self, query: np.ndarray, top_k: int = 10, nprobe: int = 64):
        """搜索最近邻"""
        assert self.trained, "索引未训练"
        if self.index_type in ("IVF", "IVF_PQ"):
            self.index.nprobe = nprobe  # 搜索的聚类数
        scores, indices = self.index.search(query.reshape(1, -1).astype(np.float32), top_k)
        return scores[0], indices[0]

    def save(self, path: str):
        faiss.write_index(self.index, path)

    def load(self, path: str):
        self.index = faiss.read_index(path)
        self.trained = True

# Milvus 使用示例
from pymilvus import connections, FieldSchema, CollectionSchema, DataType, Collection
import numpy as np

class MilvusVectorDB:
    """Milvus 向量数据库操作封装"""

    def __init__(self, host="localhost", port="19530"):
        connections.connect("default", host=host, port=port)
        self.collection = None

    def create_collection(self, name: str, dim: int, enable_metadata: bool = True):
        """创建集合"""
        fields = [
            FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
            FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=dim),
        ]

        if enable_metadata:
            fields.append(FieldSchema(name="modality", dtype=DataType.VARCHAR, max_length=20))
            fields.append(FieldSchema(name="source_id", dtype=DataType.VARCHAR, max_length=256))
            fields.append(FieldSchema(name="timestamp", dtype=DataType.INT64))

        schema = CollectionSchema(fields, description=f"Collection: {name}")
        self.collection = Collection(name, schema)

        # 创建向量索引 (IVF_FLAT)
        index_params = {
            "index_type": "IVF_FLAT",
            "metric_type": "IP",  # 内积
            "params": {"nlist": 1024}
        }
        self.collection.create_index("embedding", index_params)
        print(f"集合 {name} 创建完成, 维度={dim}")

    def insert(self, embeddings: np.ndarray, metadata: list = None):
        """插入向量数据"""
        data = [embeddings.tolist()]
        if metadata:
            data.extend([m.get(k) for m in metadata for k in metadata[0].keys()])
        self.collection.insert(data)

    def search(self, query: np.ndarray, top_k: int = 10, expr: str = ""):
        """带过滤条件的向量搜索"""
        self.collection.load()
        search_params = {"metric_type": "IP", "params": {"nprobe": 64}}

        results = self.collection.search(
            data=[query.tolist()],
            anns_field="embedding",
            param=search_params,
            limit=top_k,
            expr=expr,  # 元数据过滤条件
            output_fields=["modality", "source_id"]
        )
        return results[0]  # [Hit]

    def flush(self):
        """持久化数据"""
        self.collection.flush()

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultimodalRecommender(nn.Module):
    """三塔多模态推荐系统"""

    def __init__(self, num_users: int, num_items: int,
                 user_dim: int = 128, item_dim: int = 512,
                 text_dim: int = 768, image_dim: int = 512,
                 embed_dim: int = 128):
        super().__init__()
        # 用户塔: 用户 ID 嵌入 + 历史行为编码
        self.user_embedding = nn.Embedding(num_users, user_dim)
        self.user_tower = nn.Sequential(
            nn.Linear(user_dim, embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )

        # 物品塔: 多模态特征融合
        self.image_proj = nn.Linear(image_dim, embed_dim)
        self.text_proj = nn.Linear(text_dim, embed_dim)
        self.item_tower = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.LayerNorm(embed_dim),
            nn.ReLU(),
            nn.Linear(embed_dim, embed_dim)
        )

        # 交互层: 用户-物品匹配
        self.interaction = nn.Sequential(
            nn.Linear(embed_dim * 2, embed_dim),
            nn.ReLU(),
            nn.Dropout(0.3),
            nn.Linear(embed_dim, 1)
        )

    def forward(self, user_ids: torch.Tensor, image_feats: torch.Tensor,
                text_feats: torch.Tensor):
        # 用户表示
        user_emb = self.user_embedding(user_ids)
        user_vec = self.user_tower(user_emb)

        # 物品多模态表示
        img_vec = F.normalize(self.image_proj(image_feats), dim=-1)
        txt_vec = F.normalize(self.text_proj(text_feats), dim=-1)
        item_multimodal = torch.cat([img_vec, txt_vec], dim=-1)
        item_vec = self.item_tower(item_multimodal)

        # 交互预测
        combined = torch.cat([user_vec, item_vec], dim=-1)
        score = self.interaction(combined).squeeze(-1)
        return score

    def get_user_vector(self, user_ids: torch.Tensor):
        return self.user_tower(self.user_embedding(user_ids))

    def get_item_vector(self, image_feats: torch.Tensor, text_feats: torch.Tensor):
        img_vec = F.normalize(self.image_proj(image_feats), dim=-1)
        txt_vec = F.normalize(self.text_proj(text_feats), dim=-1)
        item_combined = torch.cat([img_vec, txt_vec], dim=-1)
        return self.item_tower(item_combined)

# 冷启动推荐：纯多模态策略
import torch
import numpy as np

class ColdStartRecommender:
    """新物品冷启动推荐：仅依赖多模态特征"""

    def __init__(self, user_vectors: np.ndarray, faiss_index):
        """
        user_vectors: 所有用户的兴趣向量 [num_users, dim]
        faiss_index: 物品多模态向量索引
        """
        self.user_vectors = user_vectors
        self.faiss_index = faiss_index

    def recommend_for_user(self, user_id: int, new_items: np.ndarray,
                          item_ids: list, top_k: int = 20):
        """为新物品找到最可能感兴趣的用户"""
        user_vec = self.user_vectors[user_id].reshape(1, -1).astype(np.float32)

        # 在多模态索引中搜索与用户兴趣最接近的新物品
        scores, indices = self.faiss_index.search(user_vec, top_k)

        results = []
        for score, idx in zip(scores[0], indices[0]):
            if idx < len(item_ids):
                results.append({
                    "item_id": item_ids[idx],
                    "match_score": float(score),
                    "reason": self._explain_match(user_id, idx)
                })
        return results

    def _explain_match(self, user_id: int, item_idx: int):
        """生成推荐解释"""
        explanations = [
            "与您收藏的商品风格相似",
            "匹配您的审美偏好",
            "与您近期浏览的商品类别一致",
            "符合您的品牌偏好"
        ]
        return explanations[item_idx % len(explanations)]

# 冷启动推荐的关键指标:
# 新物品曝光率: 提升 3-5x (相比纯协同过滤)
# 新物品点击率: 提升 2-3x
# 长尾物品覆盖率: 提升 40-60%

from fastapi import FastAPI, UploadFile, File, Form
from pydantic import BaseModel
import torch
from transformers import CLIPProcessor, CLIPModel
import faiss
import numpy as np
from PIL import Image
import io
from typing import Optional

app = FastAPI(title="图文检索系统")

# 全局模型和索引
model = None
processor = None
faiss_index = None
image_store = {}  # id -> Image

class SearchResult(BaseModel):
    image_id: str
    score: float
    description: Optional[str] = None

class SearchResponse(BaseModel):
    query: str
    results: list[SearchResult]
    total: int

@app.on_event("startup")
def load_model():
    global model, processor
    model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32")
    processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32")
    print("CLIP 模型加载完成")

@app.post("/index")
async def build_index(files: list[UploadFile] = File(...)):
    """上传图像并构建检索索引"""
    global faiss_index, image_store, model, processor

    vectors = []
    for idx, file in enumerate(files):
        content = await file.read()
        image = Image.open(io.BytesIO(content))
        inputs = processor(images=image, return_tensors="pt")
        with torch.no_grad():
            feat = model.get_image_features(**inputs)
            feat = F.normalize(feat, dim=-1)
        vectors.append(feat.numpy())
        image_store[f"img_{idx}"] = image

    vectors = np.vstack(vectors).astype(np.float32)
    dim = vectors.shape[1]
    faiss_index = faiss.IndexFlatIP(dim)
    faiss_index.add(vectors)

    return {"status": "ok", "indexed": len(files), "dimension": dim}

@app.get("/search", response_model=SearchResponse)
async def search(query: str, top_k: int = 10):
    """文本查询检索图像"""
    global model, processor, faiss_index

    inputs = processor(text=query, return_tensors="pt", padding=True)
    with torch.no_grad():
        text_feat = model.get_text_features(**inputs)
        text_feat = F.normalize(text_feat, dim=-1).numpy()

    scores, indices = faiss_index.search(text_feat, top_k)

    results = []
    for score, idx in zip(scores[0], indices[0]):
        if idx != -1:
            results.append(SearchResult(
                image_id=f"img_{idx}",
                score=float(score)
            ))

    return SearchResponse(query=query, results=results, total=len(results))

# 检索系统评估脚本
import torch
from transformers import CLIPProcessor, CLIPModel
import numpy as np
import json
from PIL import Image
from tqdm import tqdm

class RetrievalEvaluator:
    """图文检索系统评估工具"""

    def __init__(self, model_name="openai/clip-vit-base-patch32"):
        self.model = CLIPModel.from_pretrained(model_name)
        self.processor = CLIPProcessor.from_pretrained(model_name)

    def load_dataset(self, dataset_path: str):
        """加载评估数据集
        数据格式: JSONL, 每行 {"image": "path.jpg", "caption": "描述"}
        """
        self.images = []
        self.captions = []
        with open(dataset_path) as f:
            for line in f:
                data = json.loads(line)
                self.images.append(Image.open(data["image"]))
                self.captions.append(data["caption"])
        print(f"加载 {len(self.images)} 个图文对")

    @torch.no_grad()
    def evaluate(self, batch_size=64):
        """计算 Recall@K 指标"""
        # 批量编码所有图像
        image_features = []
        for i in tqdm(range(0, len(self.images), batch_size)):
            batch = self.images[i:i+batch_size]
            inputs = self.processor(images=batch, return_tensors="pt")
            feat = self.model.get_image_features(**inputs)
            image_features.append(F.normalize(feat, dim=-1))
        image_features = torch.cat(image_features, dim=0)  # [N, D]

        # 批量编码所有文本
        text_features = []
        for i in tqdm(range(0, len(self.captions), batch_size)):
            batch = self.captions[i:i+batch_size]
            inputs = self.processor(text=batch, return_tensors="pt", padding=True)
            feat = self.model.get_text_features(**inputs)
            text_features.append(F.normalize(feat, dim=-1))
        text_features = torch.cat(text_features, dim=0)  # [N, D]

        # 计算相似度矩阵
        similarity = image_features @ text_features.T  # [N, N]

        # Recall@K
        results = {}
        n = len(self.images)
        for k in [1, 5, 10, 50]:
            # 以文搜图
            _, topk_indices = similarity.topk(k, dim=0)
            t2i_correct = sum(
                1 for i in range(n) if i in topk_indices[:, i]
            )
            t2i_recall = t2i_correct / n

            # 以图搜文
            _, topk_indices = similarity.topk(k, dim=1)
            i2t_correct = sum(
                1 for i in range(n) if i in topk_indices[i]
            )
            i2t_recall = i2t_correct / n

            results[f"t2i_recall@{k}"] = f"{t2i_recall:.2%}"
            results[f"i2t_recall@{k}"] = f"{i2t_recall:.2%}"

        return results

# 典型评估结果 (MSCOCO val2017):
# t2i_recall@1:  54.2%
# t2i_recall@5:  81.7%
# t2i_recall@10: 89.3%
# i2t_recall@1:  72.4%
# i2t_recall@5:  93.1%
# i2t_recall@10: 96.5%