You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

corporations.py 3.8KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128
  1. #
  2. # Copyright 2025 The InfiniFlow Authors. All Rights Reserved.
  3. #
  4. # Licensed under the Apache License, Version 2.0 (the "License");
  5. # you may not use this file except in compliance with the License.
  6. # You may obtain a copy of the License at
  7. #
  8. # http://www.apache.org/licenses/LICENSE-2.0
  9. #
  10. # Unless required by applicable law or agreed to in writing, software
  11. # distributed under the License is distributed on an "AS IS" BASIS,
  12. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13. # See the License for the specific language governing permissions and
  14. # limitations under the License.
  15. #
  16. import logging
  17. import re
  18. import json
  19. import os
  20. import pandas as pd
  21. from rag.nlp import rag_tokenizer
  22. from . import regions
  23. current_file_path = os.path.dirname(os.path.abspath(__file__))
  24. GOODS = pd.read_csv(
  25. os.path.join(current_file_path, "res/corp_baike_len.csv"), sep="\t", header=0
  26. ).fillna(0)
  27. GOODS["cid"] = GOODS["cid"].astype(str)
  28. GOODS = GOODS.set_index(["cid"])
  29. CORP_TKS = json.load(
  30. open(os.path.join(current_file_path, "res/corp.tks.freq.json"), "r",encoding="utf-8")
  31. )
  32. GOOD_CORP = json.load(open(os.path.join(current_file_path, "res/good_corp.json"), "r",encoding="utf-8"))
  33. CORP_TAG = json.load(open(os.path.join(current_file_path, "res/corp_tag.json"), "r",encoding="utf-8"))
  34. def baike(cid, default_v=0):
  35. global GOODS
  36. try:
  37. return GOODS.loc[str(cid), "len"]
  38. except Exception:
  39. pass
  40. return default_v
  41. def corpNorm(nm, add_region=True):
  42. global CORP_TKS
  43. if not nm or not isinstance(nm, str):
  44. return ""
  45. nm = rag_tokenizer.tradi2simp(rag_tokenizer.strQ2B(nm)).lower()
  46. nm = re.sub(r"&", "&", nm)
  47. nm = re.sub(r"[\(\)()\+'\"\t \*\\【】-]+", " ", nm)
  48. nm = re.sub(
  49. r"([—-]+.*| +co\..*|corp\..*| +inc\..*| +ltd.*)", "", nm, 10000, re.IGNORECASE
  50. )
  51. nm = re.sub(
  52. r"(计算机|技术|(技术|科技|网络)*有限公司|公司|有限|研发中心|中国|总部)$",
  53. "",
  54. nm,
  55. 10000,
  56. re.IGNORECASE,
  57. )
  58. if not nm or (len(nm) < 5 and not regions.isName(nm[0:2])):
  59. return nm
  60. tks = rag_tokenizer.tokenize(nm).split()
  61. reg = [t for i, t in enumerate(tks) if regions.isName(t) and (t != "中国" or i > 0)]
  62. nm = ""
  63. for t in tks:
  64. if regions.isName(t) or t in CORP_TKS:
  65. continue
  66. if re.match(r"[0-9a-zA-Z\\,.]+", t) and re.match(r".*[0-9a-zA-Z\,.]+$", nm):
  67. nm += " "
  68. nm += t
  69. r = re.search(r"^([^a-z0-9 \(\)&]{2,})[a-z ]{4,}$", nm.strip())
  70. if r:
  71. nm = r.group(1)
  72. r = re.search(r"^([a-z ]{3,})[^a-z0-9 \(\)&]{2,}$", nm.strip())
  73. if r:
  74. nm = r.group(1)
  75. return nm.strip() + (("" if not reg else "(%s)" % reg[0]) if add_region else "")
  76. def rmNoise(n):
  77. n = re.sub(r"[\((][^()()]+[))]", "", n)
  78. n = re.sub(r"[,. &()()]+", "", n)
  79. return n
  80. GOOD_CORP = set([corpNorm(rmNoise(c), False) for c in GOOD_CORP])
  81. for c, v in CORP_TAG.items():
  82. cc = corpNorm(rmNoise(c), False)
  83. if not cc:
  84. logging.debug(c)
  85. CORP_TAG = {corpNorm(rmNoise(c), False): v for c, v in CORP_TAG.items()}
  86. def is_good(nm):
  87. global GOOD_CORP
  88. if nm.find("外派") >= 0:
  89. return False
  90. nm = rmNoise(nm)
  91. nm = corpNorm(nm, False)
  92. for n in GOOD_CORP:
  93. if re.match(r"[0-9a-zA-Z]+$", n):
  94. if n == nm:
  95. return True
  96. elif nm.find(n) >= 0:
  97. return True
  98. return False
  99. def corp_tag(nm):
  100. global CORP_TAG
  101. nm = rmNoise(nm)
  102. nm = corpNorm(nm, False)
  103. for n in CORP_TAG.keys():
  104. if re.match(r"[0-9a-zA-Z., ]+$", n):
  105. if n == nm:
  106. return CORP_TAG[n]
  107. elif nm.find(n) >= 0:
  108. if len(n) < 3 and len(nm) / len(n) >= 2:
  109. continue
  110. return CORP_TAG[n]
  111. return []