got_q_and_answer / usage.py
hash-map's picture
Upload 4 files
d4c5dff verified
import os
import json
import sentencepiece as spm
from tqdm import tqdm
MODEL_PREFIX = "icefire_spm"
sp = spm.SentencePieceProcessor(model_file=f"{MODEL_PREFIX}.model")
sample = "Daenerys Targaryen rides Drogon to Winterfell."
ids = sp.encode(sample)
pieces = sp.encode(sample, out_type=str)
print("\nTest encode:")
print(f" Text : {sample}")
print(f" IDs : {ids}")
print(f" Pieces: {pieces}")
print(f" Decode: {sp.decode(ids)}")
# check an entity
entity = "Winterfell"
e_pieces = sp.encode(entity, out_type=str)
print(f"\nEntity '{entity}' → {len(e_pieces)} piece(s): {e_pieces}")
entity = "Long Claw"
e_pieces = sp.encode(entity, out_type=str)
print(f"\nEntity '{entity}' → {len(e_pieces)} piece(s): {e_pieces}")