# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import cProfile
import pstats

def preprocess_text(text: str) -> str:
    """模拟文本预处理"""
    import re
    text = text.lower()
    text = re.sub(r"[^\w\s]", "", text)
    text = " ".join(text.split())
    return text

def process_documents(docs: list[str]) -> list[str]:
    """批量处理文档"""
    return [preprocess_text(doc) for doc in docs]

def main():
    # 模拟 10000 个文档
    docs = ["Hello, World! This is a test document." * 10] * 10_000
    result = process_documents(docs)
    return result

# 启动 profiling
profiler = cProfile.Profile()
profiler.enable()
main()
profiler.disable()

# 打印报告（按累计时间排序，只看前 15 行）
stats = pstats.Stats(profiler)
stats.sort_stats("cumulative")
stats.print_stats(15)

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    0.001    0.001    2.341    2.341 script.py:14(main)
        1    0.023    0.023    2.340    2.340 script.py:10(process_documents)
    10000    0.187    0.000    2.317    0.000 script.py:5(preprocess_text)
    10000    1.893    0.000    1.893    0.000 {built-in method re.sub}

# 直接对脚本 profiling
python -m cProfile -s cumulative your_script.py | head -30

pip install line_profiler

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

# 在需要分析的函数上加 @profile 装饰器
# 注意：运行时需要用 kernprof 命令，不是直接 python

@profile  # kernprof 会自动注入这个装饰器
def preprocess_text_v2(texts: list[str]) -> list[str]:
    import re
    result = []
    pattern = re.compile(r"[^\w\s]")    # 预编译正则
    for text in texts:
        text = text.lower()             # 第3行
        text = pattern.sub("", text)    # 第4行
        text = " ".join(text.split())   # 第5行
        result.append(text)             # 第6行
    return result

kernprof -l -v your_script.py

Line #    Hits    Time  Per Hit   % Time  Line Contents
     3   10000   45230    4.5      4.2%   text = text.lower()
     4   10000  812341   81.2     75.8%   text = pattern.sub("", text)
     5   10000  213456   21.3     19.9%   text = " ".join(text.split())

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import timeit

# 对比：普通字符串拼接 vs join
setup = """
texts = ["word"] * 1000
"""

# 方案一：+ 号拼接（每次创建新字符串，内存开销大）
stmt1 = """
result = ""
for t in texts:
    result += t + " "
"""

# 方案二：join（只创建一次字符串）
stmt2 = """
result = " ".join(texts)
"""

t1 = timeit.timeit(stmt1, setup=setup, number=1000)
t2 = timeit.timeit(stmt2, setup=setup, number=1000)
print(f"+ 号拼接: {t1:.3f}s")
print(f"join:     {t2:.3f}s")
print(f"join 快了 {t1/t2:.1f} 倍")
# 典型输出：
# + 号拼接: 0.523s
# join:     0.028s
# join 快了 18.7 倍

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import time
import numpy as np

N = 1_000_000  # 100 万次运算

# 方式一：Python 原生 for 循环
start = time.perf_counter()
result = 0.0
for i in range(N):
    result += i * 0.5
python_time = time.perf_counter() - start
print(f"Python 循环: {python_time:.3f}s")   # 约 0.08s

# 方式二：numpy 向量化
start = time.perf_counter()
arr = np.arange(N, dtype=float)
result = np.sum(arr * 0.5)
numpy_time = time.perf_counter() - start
print(f"numpy 向量化: {numpy_time:.4f}s")  # 约 0.002s

print(f"numpy 快了 {python_time/numpy_time:.0f} 倍")
# 典型输出：numpy 快了 40 倍

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import numpy as np

# 场景：计算 100 万个文本的平均词长
word_counts = np.random.randint(10, 500, size=1_000_000)  # 模拟词数
char_counts = np.random.randint(50, 5000, size=1_000_000) # 模拟字符数

# 不好的写法：Python for 循环
def avg_word_len_slow(words, chars):
    result = []
    for w, c in zip(words, chars):
        if w > 0:
            result.append(c / w)
    return result

# 好的写法：numpy 向量化
def avg_word_len_fast(words, chars):
    mask = words > 0                          # 布尔数组，表示哪些位置满足条件
    result = np.zeros(len(words))
    result[mask] = chars[mask] / words[mask]  # 只对满足条件的位置运算
    return result

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import tracemalloc

tracemalloc.start()

# 模拟可能有内存问题的代码
embeddings_cache = {}
for i in range(100_000):
    # 如果 key 是动态生成的，缓存会无限增长
    embeddings_cache[f"text_{i}"] = [0.1] * 1536  # 1536 维向量

snapshot = tracemalloc.take_snapshot()
top_stats = snapshot.statistics("lineno")

print("内存占用 TOP 10:")
for stat in top_stats[:10]:
    print(stat)

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import sys

# 列表推导式：立即计算全部，全部存入内存
list_result = [x ** 2 for x in range(1_000_000)]
print(f"列表内存: {sys.getsizeof(list_result) / 1024 / 1024:.1f} MB")
# 输出：列表内存: 8.1 MB

# 生成器表达式：不存储结果，按需计算
gen_result = (x ** 2 for x in range(1_000_000))
print(f"生成器内存: {sys.getsizeof(gen_result)} bytes")
# 输出：生成器内存: 104 bytes  ← 几乎为零！

# 在只需要遍历一次的场景下，生成器是更好的选择
# 例如：逐批把向量写入数据库，不需要全部同时在内存中
def embed_generator(texts):
    for text in texts:
        yield compute_embedding(text)  # 一次只计算一个，用完即丢

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

from openai import OpenAI
import time

client = OpenAI()
texts = ["这是一段文字"] * 1000  # 模拟 1000 条文本

# 方式一：逐条调用（1000 次 HTTP 请求）
start = time.perf_counter()
embeddings_slow = []
for text in texts:
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=text  # 每次只发一条
    )
    embeddings_slow.append(resp.data[0].embedding)
slow_time = time.perf_counter() - start
print(f"逐条处理: {slow_time:.1f}s")  # 约 50-100s

# 方式二：批量调用（分批，每批 100 条，共 10 次 HTTP 请求）
BATCH_SIZE = 100
start = time.perf_counter()
embeddings_fast = []
for i in range(0, len(texts), BATCH_SIZE):
    batch = texts[i:i + BATCH_SIZE]
    resp = client.embeddings.create(
        model="text-embedding-3-small",
        input=batch  # 一次发 100 条
    )
    embeddings_fast.extend([d.embedding for d in resp.data])
fast_time = time.perf_counter() - start
print(f"批量处理: {fast_time:.1f}s")  # 约 3-5s

print(f"批量处理快了约 {slow_time/fast_time:.0f} 倍")

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import numpy as np

# 知识库：10000 个文档，每个向量 1536 维
doc_vectors = np.random.randn(10_000, 1536).astype(np.float32)
query_vector = np.random.randn(1536).astype(np.float32)

# 方式一：Python for 循环计算余弦相似度
def cosine_sim_slow(query, docs):
    results = []
    q_norm = np.linalg.norm(query)
    for doc in docs:
        sim = np.dot(query, doc) / (q_norm * np.linalg.norm(doc))
        results.append(sim)
    return results

# 方式二：numpy 矩阵运算（一次计算全部）
def cosine_sim_fast(query, docs):
    # 归一化（每个向量除以自身的长度）
    query_norm = query / np.linalg.norm(query)
    docs_norm = docs / np.linalg.norm(docs, axis=1, keepdims=True)
    # 矩阵乘法一次计算所有相似度
    return docs_norm @ query_norm  # @ 是矩阵乘法运算符

import timeit
t1 = timeit.timeit(lambda: cosine_sim_slow(query_vector, doc_vectors), number=5)
t2 = timeit.timeit(lambda: cosine_sim_fast(query_vector, doc_vectors), number=5)
print(f"for 循环: {t1:.2f}s，numpy: {t2:.3f}s，提速 {t1/t2:.0f} 倍")
# 典型输出：for 循环: 2.43s，numpy: 0.038s，提速 64 倍

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import pandas as pd
from openai import OpenAI

client = OpenAI()

def embed_pipeline_v1(csv_path: str):
    df = pd.read_csv(csv_path)          # 全量读入内存
    results = []
    for _, row in df.iterrows():        # 逐行遍历（慢）
        text = row["content"]
        # 字符串清洗（逐条）
        text = text.strip().lower()
        # 逐条调用 Embedding API（最慢的地方）
        resp = client.embeddings.create(
            model="text-embedding-3-small",
            input=text
        )
        vector = resp.data[0].embedding
        results.append({
            "id": row["id"],
            "vector": vector
        })
    return results
# 预计耗时：500000 条 × 0.15s/条 ≈ 20 小时

# 以下为代码示例，非程序员可跳过代码，重点看文字说明

import pandas as pd
import numpy as np
from openai import OpenAI
import asyncio
import aiohttp
from typing import Generator

client = OpenAI()

# 优化一：向量化字符串清洗
def clean_texts_vectorized(df: pd.DataFrame) -> pd.DataFrame:
    df["content"] = (
        df["content"]
        .str.strip()
        .str.lower()
        .str.replace(r"\s+", " ", regex=True)
        .str[:2000]  # 截断超长文本，避免超出 token 限制
    )
    return df.dropna(subset=["content"])

# 优化二：批量 Embedding
def embed_batch(texts: list[str], batch_size: int = 256) -> list[list[float]]:
    all_embeddings = []
    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]
        resp = client.embeddings.create(
            model="text-embedding-3-small",
            input=batch
        )
        all_embeddings.extend([d.embedding for d in resp.data])
    return all_embeddings

# 优化三：分块读取大文件（控制内存峰值）
def embed_pipeline_v2(csv_path: str, output_path: str):
    CHUNK_SIZE = 10_000
    BATCH_SIZE = 256
    first_chunk = True

    for chunk in pd.read_csv(csv_path, chunksize=CHUNK_SIZE):
        # 向量化清洗
        chunk = clean_texts_vectorized(chunk)

        # 批量 Embedding
        texts = chunk["content"].tolist()
        vectors = embed_batch(texts, batch_size=BATCH_SIZE)

        # 优化四：用 numpy 转换向量为 float32（节省一半内存）
        vectors_np = np.array(vectors, dtype=np.float32)

        # 写入结果
        chunk["vector"] = vectors_np.tolist()
        chunk[["id", "vector"]].to_csv(
            output_path,
            mode="w" if first_chunk else "a",
            header=first_chunk,
            index=False
        )
        first_chunk = False
        print(f"已处理 {len(chunk)} 条")

# 优化效果对比
# v1：500000 条 × 0.15s = 20小时
# v2：500000 ÷ 256批 × 0.5s/批 ≈ 16分钟（提速 75 倍）

import numpy as np
import time

# 场景：RAG 系统中，用户提问后，需要从知识库中找到最相似的文档

# 知识库向量（模拟：10000 篇文档，每个向量 1536 维）
doc_vectors = np.random.randn(10_000, 1536).astype(np.float32)
query_vector = np.random.randn(1536).astype(np.float32)

# 方式一：Python for 循环（慢）
def find_top_k_slow(query: np.ndarray, docs: np.ndarray, k: int = 5) -> list[int]:
    """逐个计算余弦相似度，Python 循环"""
    similarities = []
    q_norm = np.linalg.norm(query)

    for i, doc in enumerate(docs):
        # 余弦相似度 = 两向量点积 / (向量1的长度 * 向量2的长度)
        sim = np.dot(query, doc) / (q_norm * np.linalg.norm(doc))
        similarities.append((i, sim))

    similarities.sort(key=lambda x: x[1], reverse=True)
    return [i for i, _ in similarities[:k]]

# 方式二：numpy 矩阵运算（快）
def find_top_k_fast(query: np.ndarray, docs: np.ndarray, k: int = 5) -> list[int]:
    """用矩阵运算一次计算所有余弦相似度"""
    # 归一化查询向量
    query_norm = query / np.linalg.norm(query)

    # 归一化所有文档向量（axis=1 表示按行计算范数）
    docs_norm = docs / np.linalg.norm(docs, axis=1, keepdims=True)

    # 矩阵乘法：一次计算所有相似度（形状 (10000,)）
    # docs_norm @ query_norm 等价于对每行做点积
    similarities = docs_norm @ query_norm

    # 找 top-k 的索引
    top_k_indices = np.argpartition(similarities, -k)[-k:]  # 比 argsort 快
    return sorted(top_k_indices.tolist(), key=lambda i: similarities[i], reverse=True)


# 性能测试
start = time.perf_counter()
for _ in range(10):
    find_top_k_slow(query_vector, doc_vectors[:1000])  # 只用1000条避免太慢
slow_time = time.perf_counter() - start

start = time.perf_counter()
for _ in range(10):
    find_top_k_fast(query_vector, doc_vectors)  # 全量10000条
fast_time = time.perf_counter() - start

print(f"Python 循环 (1000条): {slow_time:.3f}s")
print(f"numpy 向量化 (10000条): {fast_time:.3f}s")
# 典型输出：
# Python 循环 (1000条): 0.823s
# numpy 向量化 (10000条): 0.042s  ← 数据量10倍，速度却快20倍

import numpy as np
from openai import OpenAI
import time

client = OpenAI()

def embed_texts_optimized(texts: list[str], batch_size: int = 256) -> np.ndarray:
    """
    最优化的批量 Embedding 实现：
    1. 批量 API 调用（减少网络往返）
    2. 用 float32 存储（比 float64 节省一半内存）
    3. L2 归一化（便于后续用点积代替余弦相似度，更快）
    """
    all_embeddings = []

    for i in range(0, len(texts), batch_size):
        batch = texts[i:i + batch_size]

        # 批量调用：一次 API 请求，处理 256 条文本
        response = client.embeddings.create(
            model="text-embedding-3-small",
            input=batch,
            dimensions=512,  # 可以降维减少存储和计算量（原始 1536 维）
        )

        # 提取 embedding 向量列表
        batch_embeddings = [e.embedding for e in response.data]
        all_embeddings.extend(batch_embeddings)

        if (i + batch_size) % 1000 == 0:
            print(f"已处理 {min(i + batch_size, len(texts))}/{len(texts)} 条")

    # 转成 numpy 数组，用 float32 节省内存
    vectors = np.array(all_embeddings, dtype=np.float32)

    # L2 归一化：向量除以自身长度，使每个向量的欧式距离为 1
    # 归一化后，余弦相似度 = 简单的矩阵乘法（不需要除以范数）
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    vectors = vectors / norms

    return vectors


# 使用归一化后的向量进行检索（点积代替余弦相似度）
def search_normalized(
    query_vector: np.ndarray,
    doc_vectors: np.ndarray,
    k: int = 5
) -> list[int]:
    """
    对归一化向量，点积 = 余弦相似度，计算更快。
    """
    # 归一化查询向量
    query_norm = query_vector / np.linalg.norm(query_vector)

    # 矩阵乘法（已经归一化，不需要除以范数）
    similarities = doc_vectors @ query_norm

    # 取 top-k
    top_k_indices = np.argpartition(similarities, -k)[-k:]
    return sorted(top_k_indices.tolist(), key=lambda i: similarities[i], reverse=True)

# 内存占用估算
def estimate_vector_memory(n_docs: int, dimensions: int = 1536) -> dict:
    """估算向量库的内存需求"""
    # float32: 每个数字 4 字节
    # float64: 每个数字 8 字节（更精确，但 AI 任务通常不需要）
    bytes_per_vec_f32 = dimensions * 4
    bytes_per_vec_f64 = dimensions * 8

    total_f32_mb = n_docs * bytes_per_vec_f32 / (1024 * 1024)
    total_f64_mb = n_docs * bytes_per_vec_f64 / (1024 * 1024)

    return {
        "n_docs": n_docs,
        "dimensions": dimensions,
        "float32_mb": round(total_f32_mb, 1),
        "float64_mb": round(total_f64_mb, 1),
        "recommendation": "float32" if total_f32_mb < 1024 else "考虑向量数据库（Chroma/Pinecone）"
    }

# 实际场景
print(estimate_vector_memory(10_000))   # 10000 篇文档
# {'n_docs': 10000, 'dimensions': 1536, 'float32_mb': 58.6, ...}

print(estimate_vector_memory(1_000_000, dimensions=512))  # 100万篇，降维到512
# {'n_docs': 1000000, 'dimensions': 512, 'float32_mb': 1953.1, ...}
# 建议使用向量数据库