import numpy as np

a = np.array([0.023, -0.451, 0.892, -0.113])
b = np.array([0.031, -0.438, 0.901, -0.127])

# ① 余弦相似度（Cosine Similarity）
#    → 只关心"方向"，不关心"长度"
#    → 范围 [-1, 1]，越大越相似
cos_sim = np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
# → 0.9987（极像）

# ② 欧氏距离（Euclidean Distance）
#    → 关心"绝对位置"，取值范围 [0, ∞)
#    → 越小越相似，但对向量长度敏感
euclidean = np.linalg.norm(a - b)
# → 0.018（极近）

# ③ 内积（Dot Product）
#    → 既看方向又看长度
#    → 当向量已 L2 归一化时 = 余弦相似度
dot = np.dot(a, b)
# → 0.9987（因为 a、b 都是单位向量）

# 一个真实的翻车案例
import numpy as np

# 两个语义相关的句子
short_text = "猫"
long_text = "猫是一种可爱的小型哺乳动物，属于猫科，与老虎、狮子有亲缘关系"

# 未归一化的向量（早期某些模型的行为）
short_vec_unnorm = np.array([0.3, 0.4, 0.5])      # 短文本向量更短
long_vec_unnorm = np.array([0.6, 0.8, 1.0])       # 长文本向量更长

# 欧氏距离
print(np.linalg.norm(short_vec_unnorm - long_vec_unnorm))
# → 0.84（看起来不太像...但其实方向完全一样！）

# 余弦相似度
cos = np.dot(short_vec_unnorm, long_vec_unnorm) / \
      (np.linalg.norm(short_vec_unnorm) * np.linalg.norm(long_vec_unnorm))
print(cos)
# → 1.0（完全相似！）

# 教训：短文本和长文本的向量长度天然不同，
# 如果你错误地用欧氏距离做文本检索，短文本会被边缘化

from transformers import AutoModel, AutoTokenizer
import torch

model = AutoModel.from_pretrained("bert-base-chinese")
tokenizer = AutoTokenizer.from_pretrained("bert-base-chinese")

# 同一个"银行"，不同的句子
text1 = "我去银行取钱"
text2 = "河边的银行长满了草"

inputs1 = tokenizer(text1, return_tensors="pt")
inputs2 = tokenizer(text2, return_tensors="pt")

with torch.no_grad():
    emb1 = model(**inputs1).last_hidden_state[0, 1, :]  # "银行"的嵌入
    emb2 = model(**inputs2).last_hidden_state[0, 0, :]  # "银行"的嵌入（位置不同）

cos_sim = torch.cosine_similarity(emb1.unsqueeze(0), emb2.unsqueeze(0))
print(f"两个'银行'的相似度: {cos_sim.item():.4f}")
# → ~0.3-0.5（上下文不同，向量也不同了！）

# ❌ 直接拿 BERT 做句子 Embedding（效果很差！）
# 原因：BERT 没有被"训练"来区分"这篇文档和那篇文档是否相关"

# ✅ 必须经过专门的"句子嵌入训练"
# Sentence-BERT（2019）就是为此而生的

import torch
import torch.nn.functional as F

def info_nce_loss(query_emb, positive_emb, negative_embs, temperature=0.05):
    """
    query_emb:      (batch, dim) 查询向量
    positive_emb:   (batch, dim) 正例向量（相关文档）
    negative_embs:  (batch, num_negatives, dim) 负例向量（不相关文档）
    temperature:    温度参数，越低模型越"自信"（也越容易过拟合）
    """
    batch_size = query_emb.shape[0]
    
    # 正例的相似度（希望越大越好）
    pos_sim = torch.sum(query_emb * positive_emb, dim=-1) / temperature  # (batch,)
    
    # 所有负例的相似度（希望越小越好）
    neg_sim = torch.bmm(
        query_emb.unsqueeze(1),        # (batch, 1, dim)
        negative_embs.transpose(1, 2)  # (batch, dim, num_negatives)
    ).squeeze(1) / temperature          # (batch, num_negatives)
    
    # 把正例和负例的分数拼在一起
    logits = torch.cat([pos_sim.unsqueeze(1), neg_sim], dim=1)  # (batch, 1+num_negatives)
    
    # 交叉熵：正确答案是第 0 个（正例）
    labels = torch.zeros(batch_size, dtype=torch.long)
    loss = F.cross_entropy(logits, labels)
    
    return loss

# 直白理解：
# 模型在做一个"N 选 1"的选择题——
# "这些候选中，哪个是正确答案？"
# N = 1(正例) + N_neg(负例)
# 温度越低 → softmax 越尖锐 → 惩罚越严厉

"""
三类负例：Easy Negatives / Hard Negatives / False Negatives
"""

# ─── Easy Negative（简单负例）───
# 完全无关的内容，模型一眼就能区分
query = "Python 中的装饰器怎么用？"
easy_neg = "今天天气真好，适合出去散步"
# 模型："这俩差了十万八千里，轻松区分 ✅"
# → 训练价值低，模型学不到细粒度的语义边界

# ─── Hard Negative（困难负例）───
# 看起来相关但不正确，模型需要"仔细分辨"
query = "Python 中的装饰器怎么用？"
hard_neg = "Python 中如何给函数添加日志记录功能"
# 模型："都是有关于 Python 函数的，但一个问装饰器语法，一个是日志实践..."
#     "我需要仔细分辨！🤔"
# → 训练价值极高，这是模型学到"精细语义"的关键

# ─── False Negative（假负例）───
# 实际是相关/正例，但被错误标记为负例
query = "Python 中的装饰器怎么用？"
false_neg = "使用 @ 符号可以实现 Python 函数装饰器模式"
# 模型被强行告诉 "不相关"，学到的信号是错的
# → 这是毒药！降低模型精度！

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-m3")

# ─── Dense Embedding ───
dense = model.encode(
    "什么是 embedding？",
    normalize_embeddings=True
)
print(f"Dense 维度: {dense.shape}")  # (1024,)

# ─── Sparse Embedding（关键词风格）───
sparse = model.encode(
    "什么是 embedding？",
    return_dense=False,
    return_sparse=True
)
# 输出: {"什么": 0.8, "embedding": 1.2, "是": 0.3}
# 就像自动生成的 TF-IDF，但更聪明

# ─── ColBERT 风格多向量 ───
colbert_vecs = model.encode(
    "什么是 embedding？",
    return_dense=False,
    return_colbert_vecs=True
)
print(f"ColBERT 向量: {colbert_vecs.shape}")  # (N_tokens, 1024)
# 每个 token 一个向量，用于精细匹配

# GTE-Qwen2-7B 的正确使用方式
from transformers import AutoModel, AutoTokenizer
import torch

model = AutoModel.from_pretrained(
    "Alibaba-NLP/gte-Qwen2-7B-instruct",
    torch_dtype=torch.float16,
    device_map="auto",
    trust_remote_code=True
)
tokenizer = AutoTokenizer.from_pretrained(
    "Alibaba-NLP/gte-Qwen2-7B-instruct",
    trust_remote_code=True
)

# ⚠️ 必须加 Instruct Prefix！
text = "Instruct: 给定一段文本，请用向量表示其核心语义\nQuery: 中国古建筑的主要特征是什么？"

inputs = tokenizer(text, return_tensors="pt", max_length=8192, truncation=True)
with torch.no_grad():
    outputs = model(**inputs)
    # 取最后一个 token 的 hidden state（不是 mean pooling！）
    embedding = outputs.last_hidden_state[:, -1, :]
    embedding = torch.nn.functional.normalize(embedding, p=2, dim=-1)

print(f"向量维度: {embedding.shape}")  # (1, 3584)

import requests

# Jina Embeddings v3 的 task-specific API
url = "https://api.jina.ai/v1/embeddings"
headers = {"Authorization": "Bearer YOUR_JINA_API_KEY"}

# 同一个文本，不同任务 → 不同的向量！
body = {
    "model": "jina-embeddings-v3",
    "input": ["机器学习的基础概念"],
    "task": "retrieval.query",     # ← 切换这里的 task
    "dimensions": 1024
}

# 1. 用明确支持中文的模型（不要碰纯英文模型）
# ✅ BGE-M3, GTE-Qwen2-7B, text2vec-large-chinese
# ❌ all-MiniLM-L6-v2（只懂英文）

# 2. 如果中英混合场景，用 BGE-M3
# BGE-M3 在 C-MTP 数据集中包含了大量中英混合的 QA 对

# 3. 评估时用中文专属 benchmark
# C-MTEB（中文版 MTEB）的排名和英文版完全不同！

def matryoshka_loss(embeddings, labels, dims=[256, 512, 768, 1024]):
    """
    标准对比学习 loss + 对每个"截断子维度"都算一次 loss
    """
    total_loss = 0
    
    for dim in dims:
        # 截取向量的前 dim 维
        truncated_emb = embeddings[:, :dim]
        
        # 对截断后的向量算对比学习 loss
        loss = contrastive_loss(truncated_emb, labels)
        
        # 小维度给更大权重（因为它更难学好）
        weight = 1.0 / dim  # 维度越小，权重越大
        total_loss += weight * loss
    
    return total_loss

# 效果：
# 训练后的模型，1024 维 = 满血精度
# 截取前 256 维 ≈ 满血模型的 95% 精度
# 截取前 64 维 ≈ 满血模型的 90% 精度

# 使用 Matryoshka Embedding
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-m3")  # 内置 Matryoshka 支持

text = "Embedding 模型如何选型？"

# 满血 1024 维（最高精度）
full_emb = model.encode(text, normalize_embeddings=True)
print(f"完整向量: {full_emb.shape}")  # (1024,)

# 截取前 256 维（约 95% 精度，内存省 75%）
lite_emb = full_emb[:256]
print(f"精简向量: {lite_emb.shape}")  # (256,)

# 截取前 64 维（约 90% 精度，内存省 94%）
mini_emb = full_emb[:64]
print(f"迷你向量: {mini_emb.shape}")  # (64,)

from openai import OpenAI

client = OpenAI()

# 3072 维（默认，最贵最准）
emb_full = client.embeddings.create(
    model="text-embedding-3-large",
    input="你的文本",
    dimensions=3072  # 可以不传，默认 3072
)

# 256 维（matryoshka 截断，便宜）
emb_small = client.embeddings.create(
    model="text-embedding-3-large",
    input="你的文本",
    dimensions=256  # 价格按 token 计，更便宜
)

"""
微调 Embedding 模型的数据格式：
(query, positive_doc, negative_docs_list)
"""

training_data = [
    {
        "query": "患者出现 ST 段抬高的心电图表现，应该考虑什么诊断？",
        "positive": "急性心肌梗死的诊断标准包括心电图显示 ST 段抬高...",
        "negatives": [
            "今天天气晴朗适合出行",  # ← 太简单，训练价值低
            "ST 段改变也可能出现在心包炎患者中...",  # ← Hard negative！训练价值高
            "心电图 P 波的异常可能提示心房扩大...",  # ← Medium negative
        ]
    },
    # ... 越多越好，至少 5000 条
]

from sentence_transformers import SentenceTransformer, InputExample, losses
from torch.utils.data import DataLoader
import torch

# ─── 1. 加载基础模型 ───
model = SentenceTransformer("BAAI/bge-m3")

# ─── 2. 准备训练数据 ───
train_examples = []
for item in training_data:
    # 每个样本：(query, positive, negative1, negative2, ...)
    texts = [item["query"], item["positive"]] + item["negatives"]
    train_examples.append(InputExample(texts=texts))

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)

# ─── 3. 选择损失函数 ───
# MultipleNegativesRankingLoss：
#   把同一个 batch 内的 negative 当"其他样本的正例"来用
#   简单高效，不需要显式标注负例
train_loss = losses.MultipleNegativesRankingLoss(model)

# ─── 4. 训练 ───
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    optimizer_params={"lr": 2e-5},
    output_path="./bge-m3-finetuned",
    save_best_model=True,
    show_progress_bar=True,
    use_amp=True,  # 混合精度训练
)

# ─── 5. 加载微调后的模型 ───
finetuned_model = SentenceTransformer("./bge-m3-finetuned")

# ⚠️ 提示模板对 BGE 系列很重要！
# 查询侧需要加前缀：
query = "Represent this sentence for searching relevant passages: 如何微调 Embedding 模型？"

# 文档侧不需要前缀：
passage = "微调 Embedding 模型需要准备 (query, positive, negative) 三元组数据..."

query_emb = finetuned_model.encode(query)
passage_emb = finetuned_model.encode(passage)

from peft import LoraConfig, get_peft_model, TaskType

# 冻结基座模型，只训练 LoRA 权重
lora_config = LoraConfig(
    task_type=TaskType.FEATURE_EXTRACTION,
    r=16,                # LoRA rank（8-64，越大越强但也越慢）
    lora_alpha=32,       # 缩放参数
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj", "o_proj"]  # 针对哪些层加 LoRA
)

model = get_peft_model(base_model, lora_config)
model.print_trainable_parameters()
# → trainable params: 8.4M || all params: 7.62B || trainable%: 0.11%

# 用 8.4M 可训练参数、~12GB VRAM 就能微调 7B 模型

"""
Embedding 模型 Benchmark
在同一份中文数据集上对比 5 个主流模型的检索性能
"""
import numpy as np
import time
from sentence_transformers import SentenceTransformer

# ─── 评估数据集准备 ───
# 假设你有一份测试数据：1000 个 query，每个 query 有标注的正确答案
N_QUERIES = 1000
N_DOCS = 50000

# 模拟数据（实际使用时替换成你的真实数据）
np.random.seed(42)

# ─── 定义待评测模型 ───
models_to_test = {
    "BGE-small-zh": "BAAI/bge-small-zh-v1.5",
    "BGE-base-zh": "BAAI/bge-base-zh-v1.5",
    "BGE-M3": "BAAI/bge-m3",
    # GTE-Qwen2 需要单独加载方式，这里简化
}

results = {}

for model_name, model_path in models_to_test.items():
    print(f"\n{'='*50}")
    print(f"评测: {model_name}")
    print(f"{'='*50}")
    
    model = SentenceTransformer(model_path)
    
    # ─── 编码所有文档 ───
    t0 = time.time()
    doc_embeddings = model.encode(
        doc_texts,
        batch_size=256,
        show_progress_bar=True,
        normalize_embeddings=True
    )
    encode_time = time.time() - t0
    print(f"编码 {N_DOCS} 篇文档: {encode_time:.1f}s")
    
    # ─── 检索评估 ───
    top_1_hits = 0
    top_5_hits = 0
    top_10_hits = 0
    query_times = []
    
    for i, query in enumerate(query_texts):
        # 编码查询
        t0 = time.time()
        query_emb = model.encode(
            query,
            normalize_embeddings=True
        )
        query_times.append(time.time() - t0)
        
        # 计算相似度
        similarities = np.dot(doc_embeddings, query_emb)  # 归一化后等价于余弦
        
        # 取 Top-K
        top_k_indices = np.argsort(similarities)[::-1]
        
        # 检查正确答案是否在 Top-K 中
        correct_doc_id = ground_truth[i]
        if correct_doc_id == top_k_indices[0]:
            top_1_hits += 1
        if correct_doc_id in top_k_indices[:5]:
            top_5_hits += 1
        if correct_doc_id in top_k_indices[:10]:
            top_10_hits += 1
    
    results[model_name] = {
        "Top-1": top_1_hits / N_QUERIES * 100,
        "Top-5": top_5_hits / N_QUERIES * 100,
        "Top-10": top_10_hits / N_QUERIES * 100,
        "Avg Query Time (ms)": np.mean(query_times) * 1000,
        "Encode Time (s)": encode_time,
        "Model Size": model_path
    }

# ─── 打印结果 ───
print(f"\n{'='*80}")
print("Embedding 模型 Benchmark 结果")
print(f"{'='*80}")
print(f"{'模型':<20} {'Top-1%':>8} {'Top-5%':>8} {'Top-10%':>8} {'查询耗时':>10}")
print(f"{'-'*60}")
for model_name, metrics in results.items():
    print(f"{model_name:<20} "
          f"{metrics['Top-1']:>7.1f}% "
          f"{metrics['Top-5']:>7.1f}% "
          f"{metrics['Top-10']:>7.1f}% "
          f"{metrics['Avg Query Time (ms)']:>7.2f}ms")

print(f"\n预期典型结果（中文技术文档 5 万篇）:")
print(f"  BGE-small-zh:  Top-5 ~71%  |  ~2ms/query")
print(f"  BGE-base-zh:   Top-5 ~79%  |  ~3ms/query")
print(f"  BGE-M3:        Top-5 ~90%  |  ~6ms/query")
print(f"  GTE-Qwen2-7B:  Top-5 ~93%  |  ~25ms/query")

# ─── 成本效益分析 ───
print(f"\n{'='*80}")
print("精度 vs 成本分析")
print(f"{'='*80}")
print(f"""
┌─────────────────────┬──────────┬──────────┬──────────┐
│ 模型                 │ 精度排名  │ 速度排名  │ 性价比    │
├─────────────────────┼──────────┼──────────┼──────────┤
│ GTE-Qwen2-7B         │ 🥇 第1   │ 🐌 最慢   │ ⭐⭐⭐    │
│ BGE-M3               │ 🥈 第2   │ 🚀 快     │ ⭐⭐⭐⭐⭐ │
│ BGE-base-zh          │ 🥉 第3   │ 🚀 很快   │ ⭐⭐⭐⭐   │
│ BGE-small-zh         │ 🏅 第4   │ ⚡ 最快   │ ⭐⭐⭐    │
└─────────────────────┴──────────┴──────────┴──────────┘

建议：
  开发阶段 → BGE-M3（精度和速度的最佳平衡）
  生产环境 → 有 GPU → GTE-Qwen2-7B | 无 GPU → BGE-M3
  移动端/轻量场景 → BGE-small-zh
""")