| """Whole-input candidate pools for the Standard Chinese IME benchmark: Rime first, then a |
| word lattice over the rime-ice dictionary. |
| |
| For a key sequence Rime's sentence maker returns one whole-input sentence; its other |
| candidates cover only a prefix of the input. So the pool for an item is |
| 1. Rime's whole-input candidates, in Rime order (rime_candidates.py query), then |
| 2. the k best whole-input paths of a lattice built from the same rime-ice dictionaries |
| (cn_dicts/8105, base, ext, tencent, others, as imported by rime_ice.dict.yaml), |
| deduplicated, up to --n. Every candidate records its origin, so the Rime-only recall and the |
| pool recall are reported separately, and the order without an LM is Rime's. |
| |
| Lattice model: unigram over dictionary words, score = sum log(weight / total weight). Words |
| without a code (the tencent list) get codes the way librime encodes phrases: every |
| combination of their characters' readings, using only readings with at least 5% of the |
| character's total weight (librime's kMinimalWeightProportionForWordMaking; 了 liao is not |
| used for 了解 this way, 长 chang and zhang both are). Entries without a weight (others: extra |
| readings of known words) take the word's highest weight elsewhere, else 1. |
| |
| index build <work>/lattice-index.pkl from the rime-ice checkout |
| pool <rime-input.jsonl> <rime-candidates.jsonl> <out.jsonl> --track full|abbrev |
| The full track splits keys into syllables (every split is tried: xian = xian | xi an); the |
| abbreviated track reads each key letter as the first letter of one syllable. |
| """ |
| import argparse |
| import collections |
| import heapq |
| import itertools |
| import json |
| import math |
| from pathlib import Path |
| import pickle |
|
|
| ROOT=Path(__file__).resolve().parents[2] |
| ENGINE=ROOT/'datasets/ime-standard-zh-v1/engine' |
| TABLES=['cn_dicts/8105','cn_dicts/base','cn_dicts/ext','cn_dicts/tencent','cn_dicts/others'] |
| MIN_PROPORTION=0.05 |
| PER_CODE=40 |
| MAX_WORD=8 |
| ABBREV_MAX_WORD=4 |
|
|
|
|
| def entries(path): |
| body=False |
| for line in path.read_text(encoding='utf-8').split('\n'): |
| if not body: |
| body=line.strip()=='...'; continue |
| if not line or line.startswith('#'): continue |
| yield line.rstrip('\r').split('\t') |
|
|
|
|
| def build_index(shared): |
| shared=Path(shared); rows=[] |
| for t in TABLES: |
| for parts in entries(shared/f'{t}.dict.yaml'): |
| text=parts[0] |
| code=parts[1] if len(parts)>=2 and not parts[1].isdigit() else None |
| w=parts[-1] if len(parts)>=2 and parts[-1].isdigit() else None |
| rows.append((t,text,code,int(w) if w is not None else None)) |
| char_codes=collections.defaultdict(collections.Counter) |
| for t,text,code,w in rows: |
| if t=='cn_dicts/8105' and code and len(text)==1: char_codes[text][code]+=w or 0 |
| usable={c:[k for k,v in cs.items() if v>=MIN_PROPORTION*sum(cs.values())] for c,cs in char_codes.items()} |
| best_weight=collections.Counter() |
| for t,text,code,w in rows: |
| if w is not None: best_weight[text]=max(best_weight[text],w) |
| table=collections.defaultdict(dict); unencodable=0 |
| for t,text,code,w in rows: |
| w=max(1,w if w is not None else best_weight.get(text,1)) |
| if code: codes=[tuple(code.split())] |
| elif all(c in usable for c in text): |
| codes=list(itertools.islice(itertools.product(*(usable[c] for c in text)),64)) |
| else: |
| unencodable+=1; continue |
| for cd in codes: |
| if len(cd)!=len(text): continue |
| d=table[cd]; d[text]=max(d.get(text,0),w) |
| total=sum(max(d.values()) for d in table.values()) |
| full={}; abbrev=collections.defaultdict(dict) |
| for cd,d in table.items(): |
| top=heapq.nlargest(PER_CODE,d.items(),key=lambda x:x[1]) |
| full[' '.join(cd)]=[(text,math.log(w/total)) for text,w in top] |
| if len(cd)<=ABBREV_MAX_WORD: |
| a=abbrev[''.join(s[0] for s in cd)] |
| for text,w in top: a[text]=max(a.get(text,0),w) |
| abbrev={k:[(text,math.log(w/total)) for text,w in heapq.nlargest(PER_CODE,d.items(),key=lambda x:x[1])] |
| for k,d in abbrev.items()} |
| syllables=sorted({s for cd in table for s in cd if s.isalpha() and s.isascii()}) |
| return dict(full=full,abbrev=abbrev,syllables=syllables, |
| stats=dict(rows=len(rows),codes=len(table),unencodable=unencodable,syllables=len(syllables))) |
|
|
|
|
| class Lattice: |
| def __init__(self,index): |
| self.full,self.abbrev=index['full'],index['abbrev'] |
| self.syl=set(index['syllables']); self.maxlen=max(map(len,self.syl)) |
|
|
| def edges_full(self,keys): |
| """(start, end, text, logp): words whose syllables spell keys[start:end].""" |
| n=len(keys); out=[] |
| def walk(i,j,syls): |
| if syls: |
| for text,lp in self.full.get(' '.join(syls),()): out.append((i,j,text,lp)) |
| if len(syls)==MAX_WORD: return |
| for k in range(j+1,min(n,j+self.maxlen)+1): |
| if keys[j:k] in self.syl: walk(i,k,syls+[keys[j:k]]) |
| for i in range(n): walk(i,i,[]) |
| return out |
|
|
| def edges_abbrev(self,keys): |
| out=[] |
| for i in range(len(keys)): |
| for j in range(i+1,min(len(keys),i+ABBREV_MAX_WORD)+1): |
| for text,lp in self.abbrev.get(keys[i:j],()): out.append((i,j,text,lp)) |
| return out |
|
|
| def kbest(self,keys,track,k): |
| edges=self.edges_full(keys) if track=='full' else self.edges_abbrev(keys) |
| by_start=collections.defaultdict(list) |
| for i,j,text,lp in edges: by_start[i].append((j,text,lp)) |
| best=collections.defaultdict(dict); best[0]['']=0.0 |
| for i in range(len(keys)): |
| if not best[i]: continue |
| here=heapq.nlargest(k,best[i].items(),key=lambda x:x[1]) |
| for j,text,lp in by_start[i]: |
| d=best[j] |
| for prefix,s in here: |
| t=prefix+text |
| if s+lp>d.get(t,-math.inf): d[t]=s+lp |
| return [t for t,_ in heapq.nlargest(k,best[len(keys)].items(),key=lambda x:x[1])] |
|
|
|
|
| def cmd_index(a): |
| out=Path(a.work)/'lattice-index.pkl' |
| index=build_index(a.shared) |
| out.parent.mkdir(parents=True,exist_ok=True) |
| with open(out,'wb') as fh: pickle.dump(index,fh,protocol=pickle.HIGHEST_PROTOCOL) |
| print(json.dumps(dict(out=str(out),**index['stats']))) |
|
|
|
|
| def cmd_pool(a): |
| with open(Path(a.work)/'lattice-index.pkl','rb') as fh: lat=Lattice(pickle.load(fh)) |
| rime={} |
| for line in open(a.args[1],encoding='utf-8'): |
| if line.strip(): r=json.loads(line); rime[r['id']]=r |
| with open(a.args[0],encoding='utf-8') as fi, open(a.args[2],'w',encoding='utf-8') as fo: |
| for line in fi: |
| if not line.strip(): continue |
| q=json.loads(line); r=rime[q['id']]; assert r['keys']==q['keys'] |
| pool=[dict(text=t,origin='rime') for t in r['candidates'][:a.n]] |
| seen={c['text'] for c in pool} |
| for t in lat.kbest(q['keys'],a.track,a.n): |
| if len(pool)>=a.n: break |
| if len(t)==q['n_han'] and t not in seen: pool.append(dict(text=t,origin='lattice')); seen.add(t) |
| fo.write(json.dumps(dict(id=q['id'],keys=q['keys'],track=a.track,engine=r['engine'],candidates=pool), |
| ensure_ascii=False)+'\n') |
|
|
|
|
| def main(): |
| p=argparse.ArgumentParser(description=__doc__,formatter_class=argparse.RawDescriptionHelpFormatter) |
| p.add_argument('command',choices=['index','pool']) |
| p.add_argument('args',nargs='*') |
| p.add_argument('--shared',default=str(ENGINE/'rime-ice')) |
| p.add_argument('--work',default=str(ROOT/'outputs/benchmarks/ime-standard-zh-v1/rime')) |
| p.add_argument('--track',choices=['full','abbrev'],default='full') |
| p.add_argument('--n',type=int,default=20) |
| a=p.parse_args() |
| dict(index=cmd_index,pool=cmd_pool)[a.command](a) |
|
|
|
|
| if __name__=='__main__': |
| main() |
|
|
