# pip install torch
from modeling_kenga_embed import KengaEmbed

m = KengaEmbed.from_pretrained(".")
q = m.encode_queries(["Where is the capital of Russia?"])
d = m.encode_documents(["Moscow is the capital of Russia."])
print("retrieval cosine", float((q * d).sum(-1)))
a = m.encode_sts(["The cat sits on the mat."])
b = m.encode_sts(["A cat is sitting on a mat."])
print("sts cosine", float((a * b).sum(-1)))