File size: 2,719 Bytes
c6ee1cf | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 | import torch
from main import Transformer
from tokenizer import Tokenizer
import json
from safetensors.torch import load_file
with open("config.json") as f:
config = json.load(f)
def rebuild_tokenizer(word_to_idx):
tok = Tokenizer()
tok.word_to_idx = word_to_idx
tok.idx_to_word = {
int(idx): word
for word, idx in word_to_idx.items()
}
return tok
def translate(model, src_tokenizer, tgt_tokenizer, sentence, device, max_len=50):
model.eval()
src_ids = src_tokenizer.encode(sentence)
src_tensor = torch.tensor([src_ids]).to(device)
decoder_input = torch.tensor(
[[tgt_tokenizer.word_to_idx["<bos>"]]]
).to(device)
with torch.no_grad():
for _ in range(max_len):
logits = model(src_tensor, decoder_input)
next_token = logits[:, -1, :].argmax(dim=-1).item()
decoder_input = torch.cat(
[
decoder_input,
torch.tensor([[next_token]]).to(device)
],
dim=1
)
if next_token == tgt_tokenizer.word_to_idx["<eos>"]:
break
output_ids = decoder_input[0].tolist()
return tgt_tokenizer.decode(output_ids)
def main():
device = torch.device(
"cuda" if torch.cuda.is_available() else "cpu"
)
# Load vocabularies
with open("src_vocab.json", "r", encoding="utf-8") as f:
src_vocab = json.load(f)
with open("tgt_vocab.json", "r", encoding="utf-8") as f:
tgt_vocab = json.load(f)
# Rebuild tokenizers
src_tokenizer = rebuild_tokenizer(src_vocab)
tgt_tokenizer = rebuild_tokenizer(tgt_vocab)
# Load config
with open("config.json", "r", encoding="utf-8") as f:
config = json.load(f)
# Build model
model = Transformer(
src_vocab_size=config["src_vocab_size"],
tgt_vocab_size=config["tgt_vocab_size"],
d_model=config["d_model"],
num_heads=config["num_heads"],
d_ff=config["d_ff"]
).to(device)
# Load weights
weights = load_file(
"model.safetensors"
)
model.load_state_dict(weights)
model.eval()
print("Type an English sentence (or 'quit' to exit):")
while True:
sentence = input("> ")
if sentence.lower() == "quit":
break
translation = translate(
model,
src_tokenizer,
tgt_tokenizer,
sentence,
device
)
print(f"Translation: {translation}")
if __name__ == "__main__":
main() |