|
|
import ast
|
|
|
import pandas as pd
|
|
|
import numpy as np
|
|
|
from langchain_huggingface import HuggingFaceEmbeddings
|
|
|
|
|
|
model_name = "BAAI/bge-small-zh-v1.5"
|
|
|
model_kwargs = {'device': 'cpu'}
|
|
|
encode_kwargs = {'normalize_embeddings': True} # set True to compute cosine similarity
|
|
|
|
|
|
|
|
|
|
|
|
bge_hf_embedding = HuggingFaceEmbeddings(
|
|
|
model_name=model_name,
|
|
|
model_kwargs=model_kwargs,
|
|
|
encode_kwargs=encode_kwargs
|
|
|
)
|
|
|
|
|
|
# 【2】该函数就是把数据变为向量的函数
|
|
|
def text_2_embedding(text):
|
|
|
resp = bge_hf_embedding.embed_documents(
|
|
|
[text]
|
|
|
)
|
|
|
return resp[0]
|
|
|
|
|
|
# 【1】读取原始文件中的美食评论数据，通过调用Embedding模型，得到对应的向量，并保持到新文件中
|
|
|
def embedding_2_file(source_file, output_file):
|
|
|
"""读取原始的美食评论数据，通过调用Embedding模型，得到向量，并保持到新文件中"""
|
|
|
# 步骤：1、准备数据，并读取，从第index_col个字段读取
|
|
|
# pandas 是 Python 中用于数据处理和分析的核心工具，专门解决表格类数据的各类操作需求。
|
|
|
# pandas 支持读取几乎所有常见的结构化数据格式（Excel、CSV、SQL、JSON 等），并能将处理后的数据便捷地保存为这些格式，解决了数据导入导出的基础问题。
|
|
|
df = pd.read_csv(source_file, index_col=0)
|
|
|
# 读取后提取所需要的字段
|
|
|
df = df[['Time', 'ProductId', 'UserId', 'Score', 'Summary', 'Text']]
|
|
|
|
|
|
print(df.head(2)) # 打印前两行看看
|
|
|
|
|
|
# 步骤2： 清洗数据 和 合并数据
|
|
|
df = df.dropna() # 如果有空数据，就删掉
|
|
|
# 把评论的摘要 和 内容字段 合并成 一个字段（方便后续处理），放入新字段text_content
|
|
|
df['text_content'] = 'Summary: ' + df.Summary.str.strip() + "; Text: " + df.Text.str.strip()
|
|
|
print(df.head(2)) # 打印前两行看看 ，发现确实增加一个字段text_content
|
|
|
|
|
|
# 步骤3: 对text_content向量化，存到一个新的文件中，向量单独存入embedding字段
|
|
|
df['embedding'] = df.text_content.apply(lambda x: text_2_embedding(x))
|
|
|
# 保存到新文件
|
|
|
df.to_csv(output_file)
|
|
|
|
|
|
|
|
|
def cosine_distance(a, b):
|
|
|
"""计算余弦距离"""
|
|
|
return np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b))
|
|
|
|
|
|
# 【3】input：用户输入的问题 ，embedding_file 要检索的文件， top_n 按照相似性最多检索3个
|
|
|
def search_text(input, embedding_file, top_n=3):
|
|
|
"""
|
|
|
根据用户输入的问题，进行语义检索，返回最相似的前top_n个结果
|
|
|
:param input:
|
|
|
:param top_n:
|
|
|
:return:
|
|
|
"""
|
|
|
# 读取新文件
|
|
|
df_data = pd.read_csv(embedding_file)
|
|
|
# 取文件中df_data['embedding']字段，该字段在文件中是以字符串存储的
|
|
|
# 所以你从文件读取出来这个字段也是字符串类型的啊
|
|
|
# 要把这个字符串变成向量进行后续的数学计算，保持到新字段embedding_vector
|
|
|
# 目前这个新字段是在内存中的，没有在文件里
|
|
|
df_data['embedding_vector'] = df_data['embedding'].apply(ast.literal_eval)
|
|
|
|
|
|
|
|
|
# 把输入问题转化为向量
|
|
|
input_vector = text_2_embedding(input)
|
|
|
|
|
|
|
|
|
# 内存中embedding_vector字段和input_vector进行相似度比较
|
|
|
# 按照余弦相似度比较，余弦相似度单独封装到函数 cosine_distance中
|
|
|
# 产生一个新字段 similarity
|
|
|
df_data['similarity'] = df_data.embedding_vector.apply(lambda x: cosine_distance(x, input_vector))
|
|
|
|
|
|
res = (
|
|
|
# 对similarity字段排序，降序排
|
|
|
df_data.sort_values('similarity', ascending=False)
|
|
|
# 返回topn,此案例是：返回3个
|
|
|
.head(top_n)
|
|
|
# 把text_content字段的Summary: 替换为""，把; Text: 替换为""
|
|
|
.text_content.str.replace('Summary: ', "") # text_content是字段名
|
|
|
.str.replace('; Text: ', ';')
|
|
|
)
|
|
|
|
|
|
# 把topn的结果打印，每打印一行画个虚线
|
|
|
for r in res:
|
|
|
print(r)
|
|
|
print('-' * 30)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
# 可以先单独测试函数embedding_2_file('../datas/fine_food_reviews_1k.csv', '../datas/output_embedding.csv')
|
|
|
search_text('delicious beans', './output_embedding.csv') |
