matcher.py 4.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596
  1. # -*- coding: utf-8 -*-
  2. # @Author: privacy
  3. # @Date: 2024-06-27 09:33:01
  4. # @Last Modified by: privacy
  5. # @Last Modified time: 2024-09-05 10:38:48
  6. import os
  7. os.environ['TRANSFORMERS_OFFLINE'] = '1'
  8. from typing import List
  9. import torch
  10. import numpy as np
  11. import pandas as pd
  12. from sklearn.metrics.pairwise import cosine_similarity
  13. from transformers import AutoTokenizer, AutoModel
  14. class Matcher:
  15. def __init__(self):
  16. # Load model directly
  17. # # # 加载预训练的text2vec模型和分词器
  18. self.tokenizer = AutoTokenizer.from_pretrained("GanymedeNil/text2vec-base-chinese")
  19. self.model = AutoModel.from_pretrained("GanymedeNil/text2vec-base-chinese")
  20. def TopK1(self, title: str, keywords: list, query_embedding, option_embeddings: list) -> pd.Series:
  21. """
  22. 获取相似度最高的向量
  23. Args:
  24. title: 待分类词
  25. keywords: 备选类别
  26. query_embedding: 待分类词向量
  27. option_embeddings: 备选类别向量列表
  28. Returns:
  29. 类别和相似度值
  30. """
  31. # 计算相似度
  32. similarities = [cosine_similarity([query_embedding], [embedding])[0][0] for embedding in option_embeddings]
  33. # 找到最相近的关键词
  34. most_similar_keyword = keywords[similarities.index(max(similarities))]
  35. return pd.Series([most_similar_keyword, max(similarities)])
  36. def get_embedding(self, text: str) -> np.ndarray:
  37. """
  38. 单文本转换为向量
  39. Args:
  40. text: 文本
  41. Returns:
  42. text_embedding: 文本向量
  43. """
  44. encoded_input = self.tokenizer(text, return_tensors='pt')
  45. with torch.no_grad():
  46. output = self.model(**encoded_input)
  47. text_embedding = np.mean(output.last_hidden_state.mean(dim=1).numpy(), axis=0)
  48. return text_embedding
  49. def get_embeddings(self, text_list: list) -> List[np.ndarray]:
  50. """
  51. 批量文本转换为向量
  52. Args:
  53. text_list: 批量文本
  54. Returns:
  55. text_embeddings: 文本向量列表
  56. """
  57. text_embeddings = []
  58. for text in text_list:
  59. encoded_input = self.tokenizer(text, return_tensors='pt')
  60. with torch.no_grad():
  61. output = self.model(**encoded_input)
  62. text_embeddings.append(np.mean(output.last_hidden_state.mean(dim=1).numpy(), axis=0))
  63. return text_embeddings
  64. if __name__ == '__main__':
  65. matcher = Matcher()
  66. 招标因素 = ['投标人名称', '投标文件封面、投标函签字盖章', '投标文件格式', '报价唯一', '营业执照', '安全生产许可证', '资质条件', '财务要求', '业绩要求', '人员要求', '信誉要求', '不得存在的情形', '其他要求', '投标报价', '投标内容', '工期', '工程质量', '投标有效期', '投标保证金', '权利义务', '己标价工程量清单', '技术标准和要求', '其他', '以往同类项目业绩、经验', '信用评价', '财务状况', '投标报价合理性', '施工组织设计', '无机磨石品牌及质量', '无机磨石地坪的施工工艺及质量控制', '投标关键技术、设备、部件及材料的来源及供应可靠性', '施工安全和文明施工', '组织机构及施工管理人员', '价格得分']
  67. df = pd.read_json("D:\\desktop\\三峡水利\\data\\projects\\三峡左岸及地下电站地坪整治\\投标\\湖北建新建设工程有限公司_T221100130348%2F01整本文件\\投标文件-修改版9-5-1-1-title.json")
  68. del df['bbox']
  69. keyword_embeddings = matcher.get_embeddings(招标因素)
  70. result = df['text'].apply(lambda x: matcher.TopK1(x, 招标因素, matcher.get_embedding(x), keyword_embeddings))
  71. result.columns = ['因素', '相似度']
  72. df['因素'] = result['因素']
  73. df['相似度'] = result['相似度']
  74. max_sim_idx = df.groupby('因素')['相似度'].idxmax()
  75. max_sim_rows = df.loc[max_sim_idx]
  76. max_sim_rows.to_json('相似度.json', orient='records', lines=True, force_ascii=False)