import torch
import numpy as np
from scipy.spatial.distance import cosine
from modelscope.hub.snapshot_download import snapshot_download
from transformers import AutoTokenizer, AutoModel
import matplotlib.pyplot as plt
import seaborn as sns
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
# 配置参数
CACHE_DIR = "D:\\modelscope\\hub"
# 定义两个模型ID
MODEL_IDS = {
    "multilingual": "sentence-transformers/paraphrase-multilingual-MiniLM-L12-v2",
    "english": "sentence-transformers/all-MiniLM-L6-v2"
}
# 设置数据类型（半精度，减少显存占用）
TORCH_DTYPE = torch.float16
# 设置设备（优先GPU）
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
def load_model_and_tokenizer(model_id, cache_dir):
    """
    加载模型和分词器
    :param model_id: 模型ID
    :param cache_dir: 缓存目录
    :return: tokenizer, model
    """
    # 下载模型到本地
    local_path = snapshot_download(model_id, cache_dir=cache_dir)
    
    # 加载分词器
    tokenizer = AutoTokenizer.from_pretrained(
        local_path,
        trust_remote_code=True
    )
    
    # 加载模型（修正：使用AutoModel而非AutoModelForCausalLM）
    model = AutoModel.from_pretrained(
        local_path,
        trust_remote_code=True,
        torch_dtype=TORCH_DTYPE,
        device_map="auto"  # 自动分配设备
    )
    
    # 设置模型为评估模式
    model.eval()
    
    return tokenizer, model
def mean_pooling(model_output, attention_mask):
    """
    均值池化函数：生成句子嵌入
    :param model_output: 模型输出
    :param attention_mask: 注意力掩码
    :return: 归一化后的句子嵌入
    """
    # 获取token嵌入
    token_embeddings = model_output[0]  # 第一个元素是last_hidden_state
    # 扩展掩码维度，匹配token嵌入
    input_mask_expanded = attention_mask.unsqueeze(-1).expand(token_embeddings.size()).to(TORCH_DTYPE)
    # 计算均值池化（忽略padding部分）
    sum_embeddings = torch.sum(token_embeddings * input_mask_expanded, dim=1)
    sum_mask = torch.clamp(input_mask_expanded.sum(1), min=1e-9)  # 避免除以0
    mean_embeddings = sum_embeddings / sum_mask
    # L2归一化
    mean_embeddings = torch.nn.functional.normalize(mean_embeddings, p=2, dim=1)
    return mean_embeddings
def get_sentence_embedding(text, tokenizer, model):
    """
    获取单句的嵌入向量
    :param text: 输入文本
    :param tokenizer: 分词器
    :param model: 模型
    :return: 句子嵌入向量（numpy数组）
    """
    # 文本编码
    encoded_input = tokenizer(
        text,
        padding=True,
        truncation=True,
        max_length=512,
        return_tensors="pt"
    ).to(DEVICE)
    
    # 前向传播（不计算梯度，提升速度）
    with torch.no_grad():
        model_output = model(**encoded_input)
    
    # 均值池化生成句子嵌入
    sentence_embedding = mean_pooling(model_output, encoded_input['attention_mask'])
    
    # 转换为numpy数组
    return sentence_embedding.cpu().numpy()[0]
def calculate_similarity(emb1, emb2):
    """
    计算两个嵌入向量的余弦相似度
    :param emb1: 向量1
    :param emb2: 向量2
    :return: 余弦相似度值
    """
    return 1 - cosine(emb1, emb2)
# ------------------------------
# 主执行流程
# ------------------------------
if __name__ == "__main__":
    # 测试文本集（包含多语言）
    test_texts = {
        "english1": "The quick brown fox jumps over the lazy dog",
        "english2": "A fast brown fox leaps over a sleepy dog",
        "english3": "I love natural language processing",
        "chinese1": "快速的棕色狐狸跳过懒狗",
        "chinese2": "一只敏捷的棕色狐狸越过一只懒狗",
        "chinese3": "我喜欢自然语言处理"
    }
    
    # 存储结果
    results = {}
    
    # 加载并测试两个模型
    for model_name, model_id in MODEL_IDS.items():
        print(f"\n=== 加载模型：{model_name} ({model_id}) ===")
        tokenizer, model = load_model_and_tokenizer(model_id, CACHE_DIR)
        
        # 计算所有文本的嵌入
        embeddings = {}
        for text_name, text in test_texts.items():
            emb = get_sentence_embedding(text, tokenizer, model)
            embeddings[text_name] = emb
            print(f"  已生成 {text_name} 的嵌入，向量长度：{len(emb)}")
        
        # 计算相似度矩阵
        text_names = list(test_texts.keys())
        similarity_matrix = np.zeros((len(text_names), len(text_names)))
        for i, name1 in enumerate(text_names):
            for j, name2 in enumerate(text_names):
                similarity_matrix[i, j] = calculate_similarity(
                    embeddings[name1],
                    embeddings[name2]
                )
        
        # 存储结果
        results[model_name] = {
            "embeddings": embeddings,
            "similarity_matrix": similarity_matrix,
            "text_names": text_names
        }
        
        # 打印相似度矩阵
        print(f"\n{model_name} 模型相似度矩阵：")
        print(similarity_matrix.round(3))
    
    # ------------------------------
    # 可视化相似度矩阵
    # ------------------------------
    fig, axes = plt.subplots(1, 2, figsize=(16, 6))
    
    # 绘制多语言模型相似度矩阵
    sns.heatmap(
        results["multilingual"]["similarity_matrix"],
        annot=True,
        fmt=".3f",
        cmap="Blues",
        xticklabels=results["multilingual"]["text_names"],
        yticklabels=results["multilingual"]["text_names"],
        ax=axes[0]
    )
    axes[0].set_title("paraphrase-multilingual-MiniLM-L12-v2 相似度矩阵", fontsize=12)
    axes[0].set_xlabel("文本", fontsize=10)
    axes[0].set_ylabel("文本", fontsize=10)
    
    # 绘制英文模型相似度矩阵
    sns.heatmap(
        results["english"]["similarity_matrix"],
        annot=True,
        fmt=".3f",
        cmap="Blues",
        xticklabels=results["english"]["text_names"],
        yticklabels=results["english"]["text_names"],
        ax=axes[1]
    )
    axes[1].set_title("all-MiniLM-L6-v2 相似度矩阵", fontsize=12)
    axes[1].set_xlabel("文本", fontsize=10)
    axes[1].set_ylabel("文本", fontsize=10)
    
    # 保存图片
    plt.tight_layout()
    plt.savefig("model_similarity_matrix.png", dpi=300, bbox_inches="tight")
    plt.show()
    
    # ------------------------------
    # 性能对比分析
    # ------------------------------
    print("\n=== 性能对比分析 ===")
    # 提取关键相似度值
    en_similarity_multilingual = calculate_similarity(
        results["multilingual"]["embeddings"]["english1"],
        results["multilingual"]["embeddings"]["english2"]
    )
    en_similarity_english = calculate_similarity(
        results["english"]["embeddings"]["english1"],
        results["english"]["embeddings"]["english2"]
    )
    cn_similarity_multilingual = calculate_similarity(
        results["multilingual"]["embeddings"]["chinese1"],
        results["multilingual"]["embeddings"]["chinese2"]
    )
    cn_similarity_english = calculate_similarity(
        results["english"]["embeddings"]["chinese1"],
        results["english"]["embeddings"]["chinese2"]
    )
    
    print(f"英文相似句相似度 - 多语言模型：{en_similarity_multilingual:.4f}")
    print(f"英文相似句相似度 - 英文模型：{en_similarity_english:.4f}")
    print(f"中文相似句相似度 - 多语言模型：{cn_similarity_multilingual:.4f}")
    print(f"中文相似句相似度 - 英文模型：{cn_similarity_english:.4f}")
    
    # 性能对比可视化
    fig, ax = plt.subplots(figsize=(10, 6))
    categories = ['英文相似句', '中文相似句']
    multilingual_scores = [en_similarity_multilingual, cn_similarity_multilingual]
    english_scores = [en_similarity_english, cn_similarity_english]
    
    x = np.arange(len(categories))
    width = 0.35
    
    ax.bar(x - width/2, multilingual_scores, width, label='paraphrase-multilingual-MiniLM-L12-v2')
    ax.bar(x + width/2, english_scores, width, label='all-MiniLM-L6-v2')
    
    ax.set_title('模型相似度性能对比', fontsize=14)
    ax.set_ylabel('余弦相似度', fontsize=12)
    ax.set_xticks(x)
    ax.set_xticklabels(categories)
    ax.legend()
    ax.set_ylim(0, 1)
    
    # 添加数值标签
    for i, v in enumerate(multilingual_scores):
        ax.text(i - width/2, v + 0.02, f'{v:.4f}', ha='center')
    for i, v in enumerate(english_scores):
        ax.text(i + width/2, v + 0.02, f'{v:.4f}', ha='center')
    
    plt.tight_layout()
    plt.savefig("model_performance_comparison.png", dpi=300, bbox_inches="tight")
    plt.show()

import time
def test_inference_speed(model_name, model_id, cache_dir, test_texts, repeat=100):
    """
    测试模型推理速度
    :param model_name: 模型名称
    :param model_id: 模型ID
    :param cache_dir: 缓存目录
    :param test_texts: 测试文本列表
    :param repeat: 重复次数
    :return: 平均推理时间（秒/句）
    """
    tokenizer, model = load_model_and_tokenizer(model_id, cache_dir)
    
    # 预热
    for text in test_texts:
        get_sentence_embedding(text, tokenizer, model)
    
    # 正式测试
    start_time = time.time()
    for _ in range(repeat):
        for text in test_texts:
            get_sentence_embedding(text, tokenizer, model)
    end_time = time.time()
    
    total_time = end_time - start_time
    avg_time_per_sentence = total_time / (repeat * len(test_texts))
    sentences_per_second = 1 / avg_time_per_sentence
    
    print(f"\n{model_name} 推理速度测试：")
    print(f"  总耗时：{total_time:.2f} 秒（{repeat}次重复）")
    print(f"  平均每句耗时：{avg_time_per_sentence:.4f} 秒")
    print(f"  每秒处理句子数：{sentences_per_second:.1f}")
    
    return avg_time_per_sentence, sentences_per_second
# 运行速度测试
if __name__ == "__main__":
    test_texts_list = list(test_texts.values())
    multilingual_speed = test_inference_speed(
        "multilingual", MODEL_IDS["multilingual"], CACHE_DIR, test_texts_list
    )
    english_speed = test_inference_speed(
        "english", MODEL_IDS["english"], CACHE_DIR, test_texts_list
    )
    
    # 速度对比可视化
    fig, ax = plt.subplots(figsize=(8, 6))
    models = ['paraphrase-multilingual-MiniLM-L12-v2', 'all-MiniLM-L6-v2']
    speeds = [multilingual_speed[1], english_speed[1]]
    
    ax.bar(models, speeds, color=['blue', 'green'])
    ax.set_title('模型推理速度对比', fontsize=14)
    ax.set_ylabel('处理速度（句子/秒）', fontsize=12)
    ax.set_ylim(0, max(speeds) * 1.1)
    
    # 添加数值标签
    for i, v in enumerate(speeds):
        ax.text(i, v + 5, f'{v:.1f}', ha='center')
    
    plt.xticks(rotation=15)
    plt.tight_layout()
    plt.savefig("model_speed_comparison.png", dpi=300, bbox_inches="tight")
    plt.show()