| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899 |
- # -*- coding: utf-8 -*-
- """知识库: 把如东诊断 (报告/裁决/交接/现场单/六层文档/skill/findings) 切块入库; 检索 = BM25 (字符 bigram + ASCII 词), 零依赖."""
- import json, math, re
- from collections import Counter
- from pathlib import Path
- def _tokens(s):
- s = str(s).lower()
- asc = re.findall(r'[a-z0-9_#.\-]{2,}', s)
- cjk = re.findall(r'[一-鿿]', s)
- return asc + [a + b for a, b in zip(cjk, cjk[1:])]
- def _chunks_md(text, source, maxlen=700):
- out, cur, head = [], [], ''
- for line in text.split('\n'):
- if line.startswith('#'):
- if cur:
- out.append((head, '\n'.join(cur)))
- cur = []
- head = line.strip('# ').strip()
- cur.append(line)
- if sum(len(x) for x in cur) > maxlen:
- out.append((head, '\n'.join(cur)))
- cur = []
- if cur:
- out.append((head, '\n'.join(cur)))
- return [dict(source=f'{source} › {h}' if h else source, text=t.strip()) for h, t in out if t.strip()]
- PUBLIC_DOC_KEYS = ('收敛报告', '交接_振动', '现场单') # 公开模式只收结论类文档; 方法/判据/裁决/skill/findings 原文不入库
- def build(cfg):
- from src.windcms.redact import PUBLIC, redact
- docs = []
- for p in cfg['knowledge_docs']:
- p = Path(p)
- if not p.exists():
- continue
- if PUBLIC and not any(k in p.name for k in PUBLIC_DOC_KEYS):
- continue
- docs += _chunks_md(p.read_text(encoding='utf-8'), p.name)
- fj = cfg['findings']
- if fj.exists() and not PUBLIC:
- f = json.load(open(fj, encoding='utf-8'))
- for k, v in f.items():
- if k == 'meta':
- continue
- txt = json.dumps(v, ensure_ascii=False)
- for i in range(0, min(len(txt), 6000), 900):
- docs.append(dict(source=f'findings.json › {k}', text=txt[i:i + 900]))
- for d in docs:
- d['text'] = redact(d['text'])
- d['tokens'] = _tokens(d['text'] + ' ' + d['source'])
- out = cfg['out'] / 'kb.json'
- out.parent.mkdir(parents=True, exist_ok=True)
- # encoding 必须给: ensure_ascii=False 写的是中文, 缺 encoding 会按系统 locale(cp936) 编 → 中文 Windows 崩
- with open(out, 'w', encoding='utf-8') as f:
- json.dump(docs, f, ensure_ascii=False)
- print(f'kb: {len(docs)} chunks → {out}')
- _KB.pop(str(cfg['out']), None)
- return docs
- _KB = {}
- def load(cfg):
- key = str(cfg['out'])
- if key not in _KB:
- p = cfg['out'] / 'kb.json'
- _KB[key] = json.load(open(p, encoding='utf-8')) if p.exists() else build(cfg)
- return _KB[key]
- def search(cfg, query, k=5):
- docs = load(cfg)
- N = len(docs)
- q = _tokens(query)
- df = Counter()
- L = []
- for d in docs:
- L.append(len(d['tokens']))
- for t in set(d['tokens']):
- df[t] += 1
- avg = sum(L) / max(N, 1)
- out = []
- for d, l in zip(docs, L):
- tf = Counter(d['tokens'])
- sc = 0.0
- for t in q:
- if t in tf:
- idf = math.log(1 + (N - df[t] + .5) / (df[t] + .5))
- sc += idf * tf[t] * 2.2 / (tf[t] + 1.2 * (0.25 + 0.75 * l / avg))
- if sc > 0:
- out.append(dict(score=sc, source=d['source'], text=d['text']))
- return sorted(out, key=lambda x: -x['score'])[:k]
|