| 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"
|
| )
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
| src_tokenizer = rebuild_tokenizer(src_vocab)
|
| tgt_tokenizer = rebuild_tokenizer(tgt_vocab)
|
|
|
|
|
|
|
| with open("config.json", "r", encoding="utf-8") as f:
|
| config = json.load(f)
|
|
|
|
|
|
|
| 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)
|
|
|
|
|
|
|
| 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() |
