omai-research/milco-x

Usage

Loading the model

from transformers import AutoModel
model = AutoModel.from_pretrained("omai-research/milco-x", trust_remote_code=True)

Encoding text (sparse tensor)

# Returns a sparse COO tensor of shape [N, vocab_size]
sparse_reps = model.encode_text(["Baltimore: The Greatest City in America", "巴尔的摩：美国最伟大的城市", "Baltimore : La plus grande ville d'Amérique"])

Encoding text (token-weight dictionary)

# Returns a list of {token: weight} dicts, sorted by weight descending
results = model.encode_text(["Baltimore : La plus grande ville d'Amérique", "巴尔的摩：美国最伟大的城市"], return_dict=True)
print(results[0])
# {'e_baltimore': 1.8021222352981567, 'e_maryland': 1.2527629137039185, 'e_city': 1.202409029006958, 'e_largest': 0.9440274834632874, 'e_biggest': 0.9287132620811462, 'e_america': 0.890972912311554, 'e_usa': 0.8649974465370178, 'e_urban': 0.7201237678527832, 'e_geography': 0.6369074583053589, 'e_cities': 0.5816917419433594, 'e_us': 0.42353010177612305, 'e_is': 0.35060185194015503 ..}
print(results[1])
# {'e_baltimore': 1.6522579193115234, 'e_city': 1.3217558860778809, 'e_greatest': 1.1192315816879272, 'e_usa': 0.9287132620811462, 'e_maryland': 0.8281423449516296, 'e_best': 0.817850649356842, 'e_biggest': 0.8074519634246826, 'e_us': 0.6793810129165649, 'e_great': 0.630689799785614, 'e_america': 0.5255602598190308 ...}

LexEcho: Dual-view encoding (pivot + source)

When source_view=True, the LexEcho head augments the pivot representation with source-language token weights. This preserves entities and terms that may not have English equivalents:

results = model.encode_text(["巴尔的摩：美国最伟大的城市"], return_dict=True, source_view=True)
# Shape: [N, en_vocab_size + m_vocab_size]
# Columns [0, en_vocab_size) are pivot (English) terms
# Columns [en_vocab_size, en_vocab_size + m_vocab_size) are source (multilingual) terms
print(results[0])
# {'m_</s>': 2.09375, 'e_baltimore': 1.6522579193115234, 'e_city': 1.3217558860778809, 'e_greatest': 1.12432062625885, 'e_usa': 0.9225212931632996, 'e_maryland': 0.8247235417366028, 'e_best': 0.8143964409828186, 'e_biggest': 0.8039615750312805, 'm_伟大的': 0.77734375, 'm_:': 0.6875, 'm_美国': 0.68359375, 'e_us': 0.6734226942062378, 'e_great': 0.6265231370925903}