# Counts tokens for the UDHR parallel texts used in https://www.booktranslator.app/blog/llm-translation-cost
# Requirements: pip install tiktoken==0.14.0
# Data: download udhr_<code>.xml files from https://github.com/behnam/unicode-udhr/tree/547ac2beee10e1804055b63811f15034432ae559/data/udhr into ./udhr/
# Usage: python llm-translation-token-count.py > llm-translation-token-ratios.csv
import unicodedata, tiktoken, xml.etree.ElementTree as ET, hashlib, csv, sys, re
langs=[('eng','English'),('spa','Spanish'),('fra','French'),('deu_1996','German'),('por_BR','Portuguese (Brazil)'),('rus','Russian'),('arb','Arabic'),('hin','Hindi'),('vie','Vietnamese'),('tha','Thai'),('kor','Korean'),('jpn','Japanese'),('cmn_hans','Chinese (Simplified)')]
encs={n:tiktoken.get_encoding(n) for n in ['o200k_base','cl100k_base']}
def text(code,form='NFC'):
    root=ET.parse(f'udhr/udhr_{code}.xml').getroot()
    parts=[]
    for el in root.iter():
        tag=el.tag.split('}')[-1]
        if tag in ('title','para','listitem'):
            t=''.join(el.itertext()).strip()
            if t and tag!='listitem': parts.append(t)
    j='\n'.join(parts)
    return j if form=='as-stored' else unicodedata.normalize(form,j)
rows=[];base={}
variants=[(c,n,'NFC') for c,n in langs]+[('vie','Vietnamese (same file before NFC; partly decomposed)','as-stored')]
for code,name,form in variants:
    t=text(code,form); raw=open(f'udhr/udhr_{code}.xml','rb').read()
    r={'language':name,'unicode_form':form,'udhr_file':f'udhr_{code}.xml','file_sha256':hashlib.sha256(raw).hexdigest(),'characters_no_spaces':len(re.sub(r'\s','',t)),'whitespace_words':len(t.split())}
    for n,e in encs.items(): r[f'{n}_tokens']=len(e.encode(t))
    rows.append(r)
eng=rows[0]
for r in rows:
    for n in encs: r[f'{n}_vs_english']=round(r[f'{n}_tokens']/eng[f'{n}_tokens'],2)
    r['o200k_tokens_per_english_word']=round(r['o200k_base_tokens']/eng['whitespace_words'],2)
w=csv.DictWriter(sys.stdout,fieldnames=list(rows[0].keys()),lineterminator='\n');w.writeheader();w.writerows(rows)