File size: 3,757 Bytes
a449848 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """AksharaLP MultiIndic reference. Code Apache-2.0; see NOTICE for component weight licences."""
from pathlib import Path
import json,gzip,math,re,unicodedata
import numpy as np
import torch
from bhc.xlit_q import load_any
from bhc.xlit_model import beam_search_fast, set_vocab
def norm(s):return unicodedata.normalize('NFC',s.replace('\u200c','').replace('\u200d',''))
V3=dict(w_nlm=.4,w_lm=.25,w_ch_s=1.5,w_agn=2,w_agr=0,nlm_margin=2,nlm_k=6,beam_s=8,lat_k=12,rec_n=3,use_lm=True,oov_pen_s=4.)
class MultiIndic:
def __init__(self,root,language='hi',sentence=False,threads=1):
self.root=Path(root);self.threads=threads;torch.set_num_threads(threads);self.cache={};self.set_language(language)
if sentence:self.set_sentence(True)
def set_language(self,language):
if language not in ['hi','bn','te','kn','ta']:raise ValueError('Unsupported language')
p=self.root/'models'/language;set_vocab(language)
if language not in self.cache:
if language=='hi':
from bhc.translit2 import BhcTranslit2
self.cache[language]=dict(hi=BhcTranslit2(str(p),threads=self.threads))
else:
cfg=json.loads((p/'params.json').read_text());lex={}
for line in gzip.decompress((p/'dict.tsv.gz').read_bytes()).decode().splitlines():
w,c=line.rsplit('\t',1);lex[w]=math.log(int(c)/cfg['dictionary_total'])
self.cache[language]=dict(model=load_any(str(p/'xlit.q4')),p=cfg,lex=lex)
self.language=language;self.active=self.cache[language];self.sentence=False;return self
def set_sentence(self,enabled=True):
if enabled and self.language!='hi':raise ValueError('Sentence mode is Hindi-only')
if enabled and 'sent' not in self.active:
from bhc.nlm import NLM
from bhc.translit3 import BhcTranslit3
p=self.root/'models/hi';nlm=NLM.load(str(p/'sentence/nlm.q8'));self.active['sent']=BhcTranslit3(str(p),threads=self.threads,nlm=nlm,**V3)
self.sentence=enabled;return self
@torch.inference_mode()
def word(self,word,k=8):
if not re.fullmatch('[A-Za-z]+',word):return [word]
if self.language=='hi':return self.active['hi'].translit_words([word],k=k)[0]
a=self.active;c=[];seen=set()
for s,sc in beam_search_fast(a['model'],[word.lower()],beam=8,nbest=8)[0]:
s=norm(s)
if s and s not in seen:seen.add(s);c.append((s,sc))
if not c:return []
lex=a['lex'];r=a['p']['reranker'];mu=max((lex[s] for s,_ in c if s in lex),default=0);msc=max(sc for _,sc in c);out=[]
for i,(s,sc) in enumerate(c):
u=lex.get(s);inside=u is not None;f=[sc,sc/(len(s)+1),int(inside),u if inside else 0,int(i==0),i/len(c),sc-msc,u-mu if inside else 0]
score=sum((x-m)/sd*w for x,m,sd,w in zip(f,r['mu'],r['sd'],r['w']));out.append((s,score))
return [s for s,_ in sorted(out,key=lambda v:-v[1])[:k]]
def transliterate(self,text):
if self.language=='hi':
if self.sentence:return self.active['sent'].convert(text)
old=self.active['hi'].p['use_lm'];self.active['hi'].p['use_lm']=False
try:return self.active['hi'].convert(text)
finally:self.active['hi'].p['use_lm']=old
return re.sub(r'(https?://\S+|www\.\S+|[\w.+-]+@[\w.-]+\.[A-Za-z]+|[@#][\w]+)|[A-Za-z]+',lambda m:m[0] if m[1] else (self.word(m[0],1)or[m[0]])[0],text)
def load_word_gguf(self,path):
from gguf_read import load_word_model
m=load_word_model(path)
if m.cfg.get('vocab','hi')!=self.language:raise ValueError('GGUF language does not match selected language')
if self.language=='hi':
self.active['hi'].model=m
if 'sent' in self.active:self.active['sent'].model=m
else:self.active['model']=m
return self
def load_sentence_gguf(self,path):
if self.language!='hi':raise ValueError('Sentence mode is Hindi-only')
from gguf_read import load_sentence_lm
self.set_sentence(True);self.active['sent'].nlm=load_sentence_lm(path);return self
|