| from pathlib import Path
|
| from model import build_transformer
|
| from util import create_resources
|
| import torch
|
| import sys
|
| import yaml
|
|
|
| def translate(sentence: str):
|
|
|
| with open("config.yaml", "r") as file:
|
| config = yaml.safe_load(file)
|
|
|
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
| print("Using device:", device)
|
| train_dataloader,valid_dataloader,test_dataloader,tokenizer_src,tokenizer_tgt = create_resources()
|
|
|
| src_vocab_size = tokenizer_src.get_vocab_size()
|
| tgt_vocab_size = tokenizer_tgt.get_vocab_size()
|
|
|
|
|
| model = build_transformer(
|
| src_vocab_size,
|
| tgt_vocab_size,
|
| config["seq_len"],
|
| config["seq_len"],
|
| config["num_enc_dec_blocks"],
|
| config["num_of_heads"],
|
| config["d_model"]
|
| )
|
|
|
| model = model.to(device)
|
| model_filename = "models/model.pth"
|
| state = torch.load(model_filename)
|
| model.load_state_dict(state['model_state_dict'])
|
|
|
| model.eval()
|
| with torch.no_grad():
|
| source = tokenizer_src.encode(sentence)
|
| print(source,source.ids)
|
|
|
| source = torch.cat([
|
| torch.tensor([tokenizer_src.token_to_id('[SOS]')], dtype=torch.int64),
|
| torch.tensor(source.ids, dtype=torch.int64),
|
| torch.tensor([tokenizer_src.token_to_id('[EOS]')], dtype=torch.int64),
|
| torch.tensor([tokenizer_src.token_to_id('[PAD]')] * (config["seq_len"] - len(source.ids) - 2), dtype=torch.int64)
|
| ], dim=0)
|
|
|
| source = source.to(device)
|
| source = source.unsqueeze(0)
|
|
|
| print(source.shape)
|
|
|
|
|
|
|
| source_mask = (source != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
|
| print(source_mask)
|
|
|
| encoder_output = model.encode(source, source_mask)
|
| print(encoder_output.shape)
|
| decoder_input = torch.full((1, 1), tokenizer_tgt.token_to_id('[SOS]'),
|
| dtype=torch.long, device=device)
|
|
|
|
|
| while decoder_input.size(1) < config["seq_len"]:
|
|
|
| decoder_mask = torch.triu(torch.ones((1, decoder_input.size(1), decoder_input.size(1))),
|
| diagonal=1).to(device, dtype=torch.int)
|
|
|
| print("#######################")
|
|
|
|
|
| print(decoder_mask.shape)
|
| out = model.decode(decoder_input, encoder_output, source_mask, decoder_mask)
|
| prob = model.project(out[:, -1])
|
| print(max(prob[0]))
|
| print(min(prob[0]))
|
| _, next_word = torch.max(prob, dim=1)
|
|
|
| next_token = torch.full((1, 1), next_word.item(), dtype=torch.long, device=device)
|
| decoder_input = torch.cat([decoder_input, next_token], dim=1)
|
|
|
| print(f"{tokenizer_tgt.decode([next_word.item()])}", end=' ')
|
|
|
| if next_word.item() == tokenizer_tgt.token_to_id('[EOS]'):
|
| break
|
|
|
|
|
|
|
| return 0
|
|
|
|
|
| a = translate("My Name is sai and I love computers")
|
| print(a) |
