🔌 向量嵌入 Embedding

# 使用 Gensim 实现 Word2Vec
from gensim.models import Word2Vec

# 准备语料
sentences = [
    ['the', 'cat', 'sat', 'on', 'the', 'mat'],
    ['the', 'dog', 'ran', 'in', 'the', 'park'],
    ['the', 'bird', 'flew', 'over', 'the', 'tree']
]

# 训练 Word2Vec 模型
model = Word2Vec(
    sentences=sentences,
    vector_size=100,      # 向量维度
    window=5,             # 上下文窗口大小
    min_count=1,          # 最小词频
    workers=4,            # 并行线程数
    sg=1,                 # 1=Skip-gram, 0=CBOW
    epochs=10             # 训练轮数
)

# 获取词向量
vector = model.wv['cat']
print(f"'cat' 的向量：{vector[:10]}...")  # 显示前 10 维

# 查找相似词
similar = model.wv.most_similar('cat', topn=3)
print(f"与'cat'相似的词：{similar}")

# 词向量运算
result = model.wv.most_similar(positive=['king', 'woman'], negative=['man'], topn=1)
print(f"king - man + woman ≈ {result[0][0]}")

# 计算相似度
similarity = model.wv.similarity('cat', 'dog')
print(f"'cat'和'dog'的相似度：{similarity:.4f}")

# 使用 Hugging Face Transformers 获取 BERT Embedding
from transformers import BertTokenizer, BertModel
import torch

# 加载预训练模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')

# 准备文本
texts = [
    "The cat sat on the mat.",
    "The dog ran in the park."
]

# 分词并获取输入
inputs = tokenizer(texts, padding=True, truncation=True, 
                   return_tensors='pt', max_length=512)

# 获取 BERT 输出
with torch.no_grad():
    outputs = model(**inputs)

# last_hidden_state: [batch_size, seq_len, hidden_size]
last_hidden = outputs.last_hidden_state
# pooler_output: [batch_size, hidden_size] - [CLS] token 的输出
pooler_output = outputs.pooler_output

print(f"Last hidden state shape: {last_hidden.shape}")
print(f"Pooler output shape: {pooler_output.shape}")

# 获取 [CLS] token 的嵌入（常用于句子表示）
cls_embeddings = last_hidden[:, 0, :]
print(f"CLS embeddings shape: {cls_embeddings.shape}")

# 计算句子相似度
from torch.nn.functional import cosine_similarity

similarity = cosine_similarity(
    cls_embeddings[0].unsqueeze(0),
    cls_embeddings[1].unsqueeze(0)
)
print(f"句子相似度：{similarity.item():.4f}")

# 使用句向量进行语义搜索
def semantic_search(query, candidates, model, tokenizer, top_k=3):
    """语义搜索示例"""
    all_texts = [query] + candidates
    inputs = tokenizer(all_texts, padding=True, truncation=True, 
                       return_tensors='pt', max_length=512)
    
    with torch.no_grad():
        outputs = model(**inputs)
    
    embeddings = outputs.pooler_output
    query_emb = embeddings[0:1]
    candidate_embs = embeddings[1:]
    
    # 计算余弦相似度
    similarities = cosine_similarity(query_emb, candidate_embs)
    
    # 返回 top_k 结果
    top_indices = similarities.topk(top_k).indices.tolist()
    return [(candidates[i], similarities[0][i].item()) for i in top_indices]

# 测试
candidates = [
    "A feline is resting on the carpet.",
    "The canine is playing outside.",
    "Birds are flying in the sky."
]

results = semantic_search("The cat sat on the mat.", candidates, model, tokenizer)
print("\n语义搜索结果:")
for text, score in results:
    print(f"  {score:.4f}: {text}")