knowledge.py 3.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899
  1. # -*- coding: utf-8 -*-
  2. """知识库: 把如东诊断 (报告/裁决/交接/现场单/六层文档/skill/findings) 切块入库; 检索 = BM25 (字符 bigram + ASCII 词), 零依赖."""
  3. import json, math, re
  4. from collections import Counter
  5. from pathlib import Path
  6. def _tokens(s):
  7. s = str(s).lower()
  8. asc = re.findall(r'[a-z0-9_#.\-]{2,}', s)
  9. cjk = re.findall(r'[一-鿿]', s)
  10. return asc + [a + b for a, b in zip(cjk, cjk[1:])]
  11. def _chunks_md(text, source, maxlen=700):
  12. out, cur, head = [], [], ''
  13. for line in text.split('\n'):
  14. if line.startswith('#'):
  15. if cur:
  16. out.append((head, '\n'.join(cur)))
  17. cur = []
  18. head = line.strip('# ').strip()
  19. cur.append(line)
  20. if sum(len(x) for x in cur) > maxlen:
  21. out.append((head, '\n'.join(cur)))
  22. cur = []
  23. if cur:
  24. out.append((head, '\n'.join(cur)))
  25. return [dict(source=f'{source} › {h}' if h else source, text=t.strip()) for h, t in out if t.strip()]
  26. PUBLIC_DOC_KEYS = ('收敛报告', '交接_振动', '现场单') # 公开模式只收结论类文档; 方法/判据/裁决/skill/findings 原文不入库
  27. def build(cfg):
  28. from src.windcms.redact import PUBLIC, redact
  29. docs = []
  30. for p in cfg['knowledge_docs']:
  31. p = Path(p)
  32. if not p.exists():
  33. continue
  34. if PUBLIC and not any(k in p.name for k in PUBLIC_DOC_KEYS):
  35. continue
  36. docs += _chunks_md(p.read_text(encoding='utf-8'), p.name)
  37. fj = cfg['findings']
  38. if fj.exists() and not PUBLIC:
  39. f = json.load(open(fj, encoding='utf-8'))
  40. for k, v in f.items():
  41. if k == 'meta':
  42. continue
  43. txt = json.dumps(v, ensure_ascii=False)
  44. for i in range(0, min(len(txt), 6000), 900):
  45. docs.append(dict(source=f'findings.json › {k}', text=txt[i:i + 900]))
  46. for d in docs:
  47. d['text'] = redact(d['text'])
  48. d['tokens'] = _tokens(d['text'] + ' ' + d['source'])
  49. out = cfg['out'] / 'kb.json'
  50. out.parent.mkdir(parents=True, exist_ok=True)
  51. # encoding 必须给: ensure_ascii=False 写的是中文, 缺 encoding 会按系统 locale(cp936) 编 → 中文 Windows 崩
  52. with open(out, 'w', encoding='utf-8') as f:
  53. json.dump(docs, f, ensure_ascii=False)
  54. print(f'kb: {len(docs)} chunks → {out}')
  55. _KB.pop(str(cfg['out']), None)
  56. return docs
  57. _KB = {}
  58. def load(cfg):
  59. key = str(cfg['out'])
  60. if key not in _KB:
  61. p = cfg['out'] / 'kb.json'
  62. _KB[key] = json.load(open(p, encoding='utf-8')) if p.exists() else build(cfg)
  63. return _KB[key]
  64. def search(cfg, query, k=5):
  65. docs = load(cfg)
  66. N = len(docs)
  67. q = _tokens(query)
  68. df = Counter()
  69. L = []
  70. for d in docs:
  71. L.append(len(d['tokens']))
  72. for t in set(d['tokens']):
  73. df[t] += 1
  74. avg = sum(L) / max(N, 1)
  75. out = []
  76. for d, l in zip(docs, L):
  77. tf = Counter(d['tokens'])
  78. sc = 0.0
  79. for t in q:
  80. if t in tf:
  81. idf = math.log(1 + (N - df[t] + .5) / (df[t] + .5))
  82. sc += idf * tf[t] * 2.2 / (tf[t] + 1.2 * (0.25 + 0.75 * l / avg))
  83. if sc > 0:
  84. out.append(dict(score=sc, source=d['source'], text=d['text']))
  85. return sorted(out, key=lambda x: -x['score'])[:k]