File size: 2,123 Bytes
8dfcde2 | 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 | """Minimal end-to-end example: load the HF-hosted weights and translate.
Run from the parent project directory (so `src` is importable):
python example.py --text "Hello world, how are you?"
"""
import argparse
import json
from pathlib import Path
import sentencepiece as spm
import torch
# Requires: src/ from https://github.com/Euswbnix/Machine_translation on the path
from src.model import Transformer
from src.inference.translate import batched_beam_search
from src.data.tokenizer import BOS_ID, EOS_ID, PAD_ID
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--weights", default="pytorch_model.bin")
ap.add_argument("--spm", default="sentencepiece.model")
ap.add_argument("--config", default="config.json")
ap.add_argument("--text", required=True, help="English sentence to translate.")
ap.add_argument("--beam", type=int, default=5)
ap.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
args = ap.parse_args()
cfg = json.loads(Path(args.config).read_text())
model = Transformer(
vocab_size=cfg["vocab_size"],
d_model=cfg["d_model"],
n_heads=cfg["n_heads"],
n_encoder_layers=cfg["n_encoder_layers"],
n_decoder_layers=cfg["n_decoder_layers"],
d_ff=cfg["d_ff"],
dropout=0.0,
max_seq_len=cfg["max_seq_len"],
share_embeddings=cfg["share_embeddings"],
pad_idx=PAD_ID,
).to(args.device)
model.load_state_dict(torch.load(args.weights, map_location=args.device))
model.eval()
sp = spm.SentencePieceProcessor()
sp.load(args.spm)
# Encode: wrap with BOS/EOS the same way the trainer does
ids = [BOS_ID] + sp.encode(args.text, out_type=int) + [EOS_ID]
src = torch.tensor([ids], dtype=torch.long, device=args.device)
hyp_ids = batched_beam_search(
model, src, beam_size=args.beam, max_len=cfg["max_seq_len"], length_penalty=1.0
)[0]
# Strip BOS/EOS and decode
hyp_ids = [t for t in hyp_ids if t not in (BOS_ID, EOS_ID, PAD_ID)]
print(sp.decode(hyp_ids))
if __name__ == "__main__":
main()
|