# 词袋模型示例
documents = ["我爱苹果手机", "iPhone 很好用", "苹果是一种水果"]
vocab = {"我": 0, "爱": 1, "苹果": 2, "手机": 3, "iPhone": 4, "很": 5, "好用": 6, "是": 7, "一种": 8, "水果": 9}

# 文档向量（词频）
doc1 = [1, 1, 1, 1, 0, 0, 0, 0, 0, 0]  # "我爱苹果手机"
doc2 = [0, 0, 0, 0, 1, 1, 1, 0, 0, 0]  # "iPhone 很好用"
doc3 = [0, 0, 1, 0, 0, 0, 0, 1, 1, 1]  # "苹果是一种水果"

# 问题：doc1 和 doc2 在语义上相近（都是说苹果手机好），
# 但向量没有任何重叠，余弦相似度为 0

# 稠密向量示例（实际是 1536 维，这里用 4 维简化说明）
# 这些数字是模型学到的，不需要人工设计
doc1_embedding = [0.82, -0.31,  0.15,  0.64]  # "苹果手机"
doc2_embedding = [0.79, -0.28,  0.18,  0.61]  # "iPhone"，和上面很接近
doc3_embedding = [0.12,  0.67, -0.45,  0.23]  # "苹果水果"，和上面差距大

import numpy as np
from typing import Union

def cosine_similarity(a: np.ndarray, b: np.ndarray) -> float:
    """
    余弦相似度：计算两个向量夹角的余弦值
    范围 [-1, 1]，1 表示完全相同，0 表示正交（无关），-1 表示相反
    """
    # np.dot：向量点积
    # np.linalg.norm：向量的 L2 范数（模长）
    return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))

def euclidean_distance(a: np.ndarray, b: np.ndarray) -> float:
    """欧氏距离：两个向量在空间中的直线距离，越小越相似"""
    return np.linalg.norm(a - b)

# 示例：比较几个句子的语义相似度
sentences = [
    "今天天气很好",
    "天气不错，阳光明媚",
    "人工智能发展很快",
    "AI 技术进步迅猛",
]

# 向量化后（实际需要调用 Embedding API）
v1 = np.array([0.9, 0.1, 0.05, 0.02])  # "今天天气很好"
v2 = np.array([0.88, 0.12, 0.04, 0.03])  # "天气不错，阳光明媚"
v3 = np.array([0.05, 0.08, 0.91, 0.15])  # "人工智能发展很快"
v4 = np.array([0.04, 0.07, 0.89, 0.14])  # "AI 技术进步迅猛"

print(cosine_similarity(v1, v2))  # ~0.998，非常相似（同一话题）
print(cosine_similarity(v1, v3))  # ~0.15，很不相似（不同话题）
print(cosine_similarity(v3, v4))  # ~0.999，非常相似（都是讲 AI）

from openai import OpenAI
import numpy as np

client = OpenAI()

def get_embedding(text: str, model: str = "text-embedding-3-small") -> np.ndarray:
    """
    获取单条文本的 Embedding 向量
    text-embedding-3-small：1536 维，每百万 token $0.02（便宜但足够用）
    text-embedding-3-large：3072 维，每百万 token $0.13（精度更高）
    """
    # 清理文本：换行符会影响 Embedding 质量
    text = text.replace("\n", " ").strip()
    response = client.embeddings.create(input=text, model=model)
    return np.array(response.data[0].embedding)

def get_embeddings_batch(
    texts: list[str],
    model: str = "text-embedding-3-small"
) -> list[np.ndarray]:
    """
    批量获取 Embedding，一次 API 调用处理多条文本
    比逐条调用快很多，且节省 API 开销
    """
    texts = [t.replace("\n", " ").strip() for t in texts]
    response = client.embeddings.create(input=texts, model=model)
    # 按 index 排序，确保顺序和输入一致
    sorted_data = sorted(response.data, key=lambda x: x.index)
    return [np.array(item.embedding) for item in sorted_data]

# 计算两段文本的语义相似度
def semantic_similarity(text1: str, text2: str) -> float:
    """计算两段文本的语义相似度，返回 0 到 1 之间的分数"""
    emb1, emb2 = get_embeddings_batch([text1, text2])
    # 余弦相似度可能是负数，这里映射到 [0, 1]
    cosine = np.dot(emb1, emb2) / (np.linalg.norm(emb1) * np.linalg.norm(emb2))
    return float((cosine + 1) / 2)

# 示例
score = semantic_similarity("苹果手机的价格", "iPhone 多少钱")
print(f"相似度：{score:.3f}")  # 约 0.95，非常相似

# 指定输出 256 维（而不是默认的 1536 维）
# 存储空间减少 6 倍，相似度精度小幅下降
response = client.embeddings.create(
    input="文本内容",
    model="text-embedding-3-small",
    dimensions=256,  # 截断到 256 维
)

from sentence_transformers import SentenceTransformer
import numpy as np

# 加载 BGE 模型（首次会下载约 600MB）
model = SentenceTransformer("BAAI/bge-large-zh-v1.5")

def get_chinese_embeddings(texts: list[str]) -> np.ndarray:
    """
    使用 BGE 中文模型获取 Embedding
    BGE 模型建议在查询文本前加 "为这个句子生成表示以用于检索相关文章："
    这是 BGE 的特定训练指令，能提升检索效果
    """
    # 对查询文本加前缀（对文档文本不需要加）
    texts_with_instruction = [
        f"为这个句子生成表示以用于检索相关文章：{t}" for t in texts
    ]
    embeddings = model.encode(texts_with_instruction, normalize_embeddings=True)
    return embeddings

import numpy as np
from dataclasses import dataclass
from openai import OpenAI

client = OpenAI()

@dataclass
class Document:
    id: str
    content: str
    embedding: np.ndarray | None = None

class SimpleSemanticSearch:
    """
    基于 Embedding 的简单语义搜索
    小数据量（<10万条）用暴力搜索就够了，不需要向量数据库
    """

    def __init__(self, model: str = "text-embedding-3-small"):
        self.model     = model
        self.documents: list[Document] = []

    def add_documents(self, docs: list[dict]) -> None:
        """添加文档并生成 Embedding（批量处理）"""
        texts    = [d["content"] for d in docs]
        embeddings = self._get_embeddings(texts)
        for doc, emb in zip(docs, embeddings):
            self.documents.append(Document(
                id=doc["id"],
                content=doc["content"],
                embedding=emb,
            ))

    def search(self, query: str, top_k: int = 5) -> list[dict]:
        """
        语义搜索：返回最相似的 top_k 篇文档
        时间复杂度 O(n)，n 是文档数量
        """
        query_emb = self._get_embeddings([query])[0]

        # 计算查询向量与所有文档向量的余弦相似度
        scores = []
        for doc in self.documents:
            score = np.dot(query_emb, doc.embedding) / (
                np.linalg.norm(query_emb) * np.linalg.norm(doc.embedding)
            )
            scores.append((score, doc))

        # 按相似度降序排列，取前 top_k
        scores.sort(key=lambda x: x[0], reverse=True)
        return [
            {"id": doc.id, "content": doc.content, "score": float(score)}
            for score, doc in scores[:top_k]
        ]

    def _get_embeddings(self, texts: list[str]) -> list[np.ndarray]:
        response = client.embeddings.create(input=texts, model=self.model)
        sorted_data = sorted(response.data, key=lambda x: x.index)
        return [np.array(item.embedding) for item in sorted_data]

# 使用示例
searcher = SimpleSemanticSearch()
searcher.add_documents([
    {"id": "1", "content": "Python 是一门动态类型的编程语言"},
    {"id": "2", "content": "Java 是静态类型的面向对象语言"},
    {"id": "3", "content": "机器学习需要大量的数据"},
    {"id": "4", "content": "深度学习是机器学习的子领域"},
    {"id": "5", "content": "PyTorch 是常用的深度学习框架"},
])

results = searcher.search("AI 训练需要什么", top_k=3)
for r in results:
    print(f"[{r['score']:.3f}] {r['content']}")
# [0.912] 机器学习需要大量的数据
# [0.876] 深度学习是机器学习的子领域
# [0.834] PyTorch 是常用的深度学习框架