136 lines
5.0 KiB
Python
136 lines
5.0 KiB
Python
#!/usr/bin/env python
|
||
# -*- coding: utf-8 -*-
|
||
#
|
||
# Copyright @2024 AI. Inspur Inc.
|
||
#
|
||
# @author: suojiayi <suojiayi@inspur.com>
|
||
# @date: 2025/05/13
|
||
|
||
import torch
|
||
import pandas as pd
|
||
from torch.nn.functional import normalize
|
||
import faiss
|
||
import numpy as np
|
||
from transformers import AutoTokenizer, AutoModel
|
||
from tqdm import tqdm
|
||
from loguru import logger
|
||
|
||
# 设置随机种子
|
||
def set_seed(seed=42):
|
||
import random
|
||
import numpy as np
|
||
import torch
|
||
random.seed(seed)
|
||
np.random.seed(seed)
|
||
torch.manual_seed(seed)
|
||
if torch.cuda.is_available():
|
||
torch.cuda.manual_seed_all(seed)
|
||
|
||
# 文本去重类
|
||
class SemanticDeduplicator:
|
||
def __init__(self, model_name_or_path: str, device: str = "cuda"):
|
||
self.device = torch.device(device if torch.cuda.is_available() else "cpu")
|
||
self.tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
|
||
self.model = AutoModel.from_pretrained(model_name_or_path).to(self.device)
|
||
self.model.eval()
|
||
self.dimension = self.model.config.hidden_size
|
||
self.index = faiss.IndexFlatIP(self.dimension) # 内积即余弦相似度
|
||
self.seen_embeddings = []
|
||
|
||
@torch.no_grad()
|
||
def get_embedding(self, texts):
|
||
inputs = self.tokenizer(texts, padding=True, truncation=True, max_length=512, return_tensors="pt").to(self.device)
|
||
outputs = self.model(**inputs)
|
||
embeddings = outputs.last_hidden_state[:, 0, :] # 取 [CLS] 向量作为句子表示
|
||
embeddings = normalize(embeddings, p=2, dim=1).cpu().numpy()
|
||
return embeddings
|
||
|
||
def deduplicate(self, texts, threshold=0.85):
|
||
result = []
|
||
for text in tqdm(texts, desc="De-duplicating"):
|
||
if not text.strip():
|
||
continue
|
||
emb = self.get_embedding([text])
|
||
if len(self.seen_embeddings) > 0:
|
||
self.index.add(np.array(self.seen_embeddings))
|
||
D, I = self.index.search(emb, k=1)
|
||
if '6+6' in texts:
|
||
print(D[0][0])
|
||
if '4+4' in texts:
|
||
print(D[0][0])
|
||
if D[0][0] < threshold:
|
||
result.append(text)
|
||
self.seen_embeddings.append(emb[0])
|
||
self.index.reset() # 每次 search 后要清空 index 缓存
|
||
else:
|
||
result.append(text)
|
||
self.seen_embeddings.append(emb[0])
|
||
return result
|
||
|
||
def main():
|
||
# 配置参数
|
||
input_path = './dhbq/frequency-score-5.xlsx' # 输入Excel文件路径
|
||
output_path = './dhbq/frequency-score-5-deduplicated.xlsx' # 输出Excel文件路径
|
||
model_name_or_path = "/model-pvc/suojiayi/bge-base-zh-v1.5/" # 使用的预训练模型名称
|
||
col_index = 0 # 需要去重的列索引
|
||
threshold = 0.65 # 相似度阈值
|
||
|
||
logger.info("加载数据...")
|
||
df = pd.read_excel(input_path).fillna("")
|
||
texts = df.iloc[:, col_index].astype(str).str.strip().tolist()
|
||
|
||
logger.info(f"开始语义去重（模型: {model_name_or_path} | 相似度阈值: {threshold}）")
|
||
deduplicator = SemanticDeduplicator(model_name_or_path)
|
||
unique_texts = deduplicator.deduplicate(texts=texts, threshold=threshold)
|
||
|
||
logger.info(f"去重完成，共保留 {len(unique_texts)} 条文本")
|
||
|
||
new_df = pd.DataFrame({df.columns[col_index]: unique_texts})
|
||
|
||
merged_df = pd.merge(new_df, df, on=df.columns[col_index], how='left')
|
||
|
||
# 方法2: 使用 map（适合映射单列）
|
||
# new_df[df.columns[1]] = new_df[df.columns[col_index]].map(df.set_index(df.columns[col_index])[df.columns[1]])
|
||
# new_df[df.columns[2]] = new_df[df.columns[col_index]].map(df.set_index(df.columns[col_index])[df.columns[2]])
|
||
|
||
# 如果你只想保留 new_df 中存在的列，可以这样：
|
||
final_df = merged_df[df.columns] # 按原始顺序保留所有列
|
||
|
||
logger.info(f"保存结果到 {output_path}")
|
||
final_df.to_excel(output_path, index=False)
|
||
logger.info("保存成功！")
|
||
|
||
if __name__ == "__main__":
|
||
set_seed()
|
||
main()
|
||
|
||
|
||
#
|
||
# from transformers import AutoTokenizer, AutoModel
|
||
# import torch
|
||
# import numpy as np
|
||
#
|
||
# def get_embedding(model, tokenizer, text):
|
||
# inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt").to("cuda")
|
||
# with torch.no_grad():
|
||
# outputs = model(**inputs)
|
||
# embeddings = outputs.last_hidden_state[:, 0, :].cpu().numpy() # [CLS] 向量
|
||
# return embeddings
|
||
#
|
||
# def cos_sim(a, b):
|
||
# return np.dot(a, b.T) / (np.linalg.norm(a) * np.linalg.norm(b))
|
||
#
|
||
# # 加载模型
|
||
# model_path = "/model-pvc/suojiayi/bge-base-zh-v1.5/"
|
||
# tokenizer = AutoTokenizer.from_pretrained(model_path)
|
||
# model = AutoModel.from_pretrained(model_path).to("cuda")
|
||
#
|
||
# # 测试文本
|
||
# text1 = "今天天气"
|
||
# text2 = "今天天气怎么样"
|
||
#
|
||
# emb1 = get_embedding(model, tokenizer, text1)
|
||
# emb2 = get_embedding(model, tokenizer, text2)
|
||
#
|
||
# similarity = cos_sim(emb1, emb2)
|
||
# print(f"相似度：{similarity[0][0]:.4f}") |