LangChain之特征

1. 获取 Embedding

from langchain.embeddings import HuggingFaceEmbeddings
from transformers import AutoTokenizer, AutoModel
import torch

model_name = "sentence-transformers/all-MiniLM-L6-v2"

# 初始化分词器和模型
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModel.from_pretrained(model_name)

# 使用 LangChain 的 HuggingFaceEmbeddings 包装器
embeddings = HuggingFaceEmbeddings(model_name=model_name)

1.1. 获取单句子

text = "这是一个测试文本。"

# 获取嵌入
embedding = embeddings.embed_query(text)

print(embedding)

1.2. 获取批量

texts = ["这是一个测试文本。", "这是另一个测试文本。"]

# 使用 LangChain 的 embed_documents 方法
embeddings = embeddings.embed_documents(texts)

for text, embedding in zip(texts, embeddings):
    print(f"Text: {text}")
    print(f"Embedding: {embedding}\n")

1.3. 手动控制

def get_embedding(text, tokenizer, model):
    # 对文本进行分词
    inputs = tokenizer(text, return_tensors='pt')
    
    # 获取模型的输出
    with torch.no_grad():
        outputs = model(**inputs)
    
    # 获取嵌入向量（通常取 [CLS] 令牌的输出）
    embedding = outputs.last_hidden_state[:, 0, :].squeeze().numpy()
    
    return embedding

text = "这是一个测试文本。"
embedding = get_embedding(text, tokenizer, model)

print(embedding)