pt-sk
/

transformer_eng-it

Model card Files Files and versions Metrics Training metrics Community

transformer_eng-it / Others /translate.py

pt-sk's picture

Upload 9 files

08b8da2 verified 3 months ago

raw history blame contribute delete

No virus

3.86 kB

	from pathlib import Path
	from config import get_config, latest_weights_file_path
	from model import build_transformer
	from tokenizers import Tokenizer
	from datasets import load_dataset
	from dataset import BilingualDataset
	import torch
	import sys

	def translate(sentence: str):
	# Define the device, tokenizers, and model
	device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
	print("Using device:", device)
	config = get_config()
	tokenizer_src = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_src']))))
	tokenizer_tgt = Tokenizer.from_file(str(Path(config['tokenizer_file'].format(config['lang_tgt']))))
	model = build_transformer(tokenizer_src.get_vocab_size(), tokenizer_tgt.get_vocab_size(), config["seq_len"], config['seq_len'], d_model=config['d_model']).to(device)

	# Load the pretrained weights
	model_filename = latest_weights_file_path(config)
	state = torch.load(model_filename)
	model.load_state_dict(state['model_state_dict'])

	# if the sentence is a number use it as an index to the test set
	label = ""
	if type(sentence) == int or sentence.isdigit():
	id = int(sentence)
	ds = load_dataset(f"{config['datasource']}", f"{config['lang_src']}-{config['lang_tgt']}", split='all')
	ds = BilingualDataset(ds, tokenizer_src, tokenizer_tgt, config['lang_src'], config['lang_tgt'], config['seq_len'])
	sentence = ds[id]['src_text']
	label = ds[id]["tgt_text"]
	seq_len = config['seq_len']

	# translate the sentence
	model.eval()
	with torch.no_grad():
	# Precompute the encoder output and reuse it for every generation step
	source = tokenizer_src.encode(sentence)
	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]')] * (seq_len - len(source.ids) - 2), dtype=torch.int64)
	], dim=0).to(device)
	source_mask = (source != tokenizer_src.token_to_id('[PAD]')).unsqueeze(0).unsqueeze(0).int().to(device)
	encoder_output = model.encode(source, source_mask)

	# Initialize the decoder input with the sos token
	decoder_input = torch.empty(1, 1).fill_(tokenizer_tgt.token_to_id('[SOS]')).type_as(source).to(device)

	# Print the source sentence and target start prompt
	if label != "": print(f"{f'ID: ':>12}{id}")
	print(f"{f'SOURCE: ':>12}{sentence}")
	if label != "": print(f"{f'TARGET: ':>12}{label}")
	print(f"{f'PREDICTED: ':>12}", end='')

	# Generate the translation word by word
	while decoder_input.size(1) < seq_len:
	# build mask for target and calculate output
	decoder_mask = torch.triu(torch.ones((1, decoder_input.size(1), decoder_input.size(1))), diagonal=1).type(torch.int).type_as(source_mask).to(device)
	out = model.decode(encoder_output, source_mask, decoder_input, decoder_mask)

	# project next token
	prob = model.project(out[:, -1])
	_, next_word = torch.max(prob, dim=1)
	decoder_input = torch.cat([decoder_input, torch.empty(1, 1).type_as(source).fill_(next_word.item()).to(device)], dim=1)

	# print the translated word
	print(f"{tokenizer_tgt.decode([next_word.item()])}", end=' ')

	# break if we predict the end of sentence token
	if next_word == tokenizer_tgt.token_to_id('[EOS]'):
	break

	# convert ids to tokens
	return tokenizer_tgt.decode(decoder_input[0].tolist())

	#read sentence from argument
	translate(sys.argv[1] if len(sys.argv) > 1 else "I am not a very good a student.")