| import argparse |
| import torch |
| import sentencepiece as spm |
| import logging |
|
|
| from src.models.transformer.seq2seq import BaselineSeq2SeqTransformer |
| from src.models.transformer.tokenizer import PAD_IDX, SOS_IDX, EOS_IDX |
| from src.models.transformer.helpers import generate_square_subsequent_mask |
|
|
| from scripts.preprocess import preprocess_text |
|
|
| logging.basicConfig(level=logging.INFO) |
| logger = logging.getLogger(__name__) |
|
|
|
|
| def beam_search_decode( |
| model: torch.nn.Module, |
| src_tensor: torch.Tensor, |
| device: torch.device, |
| beam_size: int = 5, |
| max_len: int = 40, |
| alpha: float = 0.7, |
| ): |
| """Executes Beam Search decoding to find the most probable translation sequence.""" |
| model.eval() |
| src_tensor = src_tensor.to(device) |
|
|
| with torch.no_grad(): |
| src_mask = torch.zeros( |
| (src_tensor.shape[0], src_tensor.shape[0]), device=device, dtype=torch.bool |
| ) |
|
|
| with torch.amp.autocast("cuda" if torch.cuda.is_available() else "cpu"): |
| memory = model.encode(src_tensor, src_mask) |
|
|
| beams = [([SOS_IDX], 0.0)] |
| completed = [] |
|
|
| for _ in range(max_len): |
| candidates = [] |
|
|
| for seq, score in beams: |
| if seq[-1] == EOS_IDX: |
| completed.append((seq, score)) |
| continue |
|
|
| trg_tensor = torch.tensor( |
| seq, dtype=torch.long, device=device |
| ).unsqueeze(1) |
| trg_mask = generate_square_subsequent_mask(trg_tensor.size(0)).to( |
| device |
| ) |
|
|
| with torch.amp.autocast("cuda" if torch.cuda.is_available() else "cpu"): |
| out = model.decode(trg_tensor, memory, trg_mask) |
| logits = model.generator(out[-1]) |
|
|
| probs = torch.log_softmax(logits, dim=-1).squeeze(0) |
| top_probs, top_idxs = torch.topk(probs, beam_size) |
|
|
| for p, idx in zip(top_probs, top_idxs): |
| idx = idx.item() |
| p = p.item() |
|
|
| |
| if len(seq) > 2 and idx == seq[-1] and idx == seq[-2]: |
| continue |
|
|
| candidates.append((seq + [idx], score + p)) |
|
|
| beams = sorted( |
| candidates, |
| key=lambda x: x[1] / ((len(x[0]) ** alpha) + 1e-6), |
| reverse=True, |
| )[:beam_size] |
|
|
| if beams and all(seq[-1] == EOS_IDX for seq, _ in beams): |
| break |
|
|
| completed.extend(beams) |
| if not completed: |
| return beams[0][0] |
|
|
| best_seq = sorted( |
| completed, |
| key=lambda x: x[1] / ((len(x[0]) ** alpha) + 1e-6), |
| reverse=True, |
| )[0][0] |
|
|
| return best_seq |
|
|
|
|
| class TransformerTranslator: |
| """ |
| A unified wrapper class for the Transformer model. |
| Instantiate this twice in your API server (once for WAR-TGL, once for TGL-WAR). |
| """ |
|
|
| def __init__(self, model_path: str, tokenizer_path: str): |
| self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| logger.info(f"Initializing Translator API on: {self.device} using {model_path}") |
|
|
| try: |
| |
| self.sp = spm.SentencePieceProcessor() |
| self.sp.load(str(tokenizer_path)) |
| vocab_size = self.sp.get_piece_size() |
|
|
| |
| self.model = BaselineSeq2SeqTransformer( |
| num_encoder_layers=6, |
| num_decoder_layers=6, |
| emb_size=512, |
| nhead=8, |
| src_vocab_size=vocab_size, |
| tgt_vocab_size=vocab_size, |
| dim_feedforward=2048, |
| dropout=0.3, |
| ).to(self.device) |
|
|
| |
| self.model.load_state_dict( |
| torch.load(model_path, map_location=self.device, weights_only=True) |
| ) |
| self.model.eval() |
| logger.info("Model weights loaded and ready for inference.") |
|
|
| except Exception as e: |
| logger.error(f"Failed to load model or tokenizer: {str(e)}") |
| raise e |
|
|
| def predict(self, sentence: str, beam_size: int = 3) -> str: |
| """Takes a single raw string and returns the translated string.""" |
| if not sentence or not isinstance(sentence, str) or not sentence.strip(): |
| return "" |
|
|
| try: |
| |
| cleaned_sentence = preprocess_text(sentence.lower()) |
| raw_tokens = self.sp.encode(str(cleaned_sentence), out_type=int)[:100] |
| src_tokens = [SOS_IDX] + raw_tokens + [EOS_IDX] |
| src_tensor = torch.tensor(src_tokens, dtype=torch.long).unsqueeze(1) |
|
|
| |
| predicted_ids = beam_search_decode( |
| self.model, src_tensor, self.device, beam_size=beam_size |
| ) |
|
|
| clean_ids = [ |
| idx for idx in predicted_ids if idx not in [SOS_IDX, EOS_IDX, PAD_IDX] |
| ] |
| translation = self.sp.decode(clean_ids) |
|
|
| return translation |
|
|
| except Exception as e: |
| logger.error(f"Translation generation failed: {str(e)}") |
| return "[Error: Unable to process translation]" |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Single-Input Transformer Inference") |
| parser.add_argument("--text", type=str, required=True, help="The source sentence") |
| parser.add_argument( |
| "--direction", |
| type=str, |
| choices=["war-tgl", "tgl-war"], |
| required=True, |
| help="Translation direction", |
| ) |
| args = parser.parse_args() |
|
|
| |
| if args.direction == "war-tgl": |
| m_path = "models/transformer/waray_tagalog_transformer_model.pt" |
| t_path = "models/transformer/tokenizer/waray_tagalog_bpe.model" |
| else: |
| m_path = "models/transformer/tagalog_waray_transformer_model.pt" |
| t_path = "models/transformer/tokenizer/tagalog_waray_bpe.model" |
|
|
| |
| translator = TransformerTranslator(model_path=m_path, tokenizer_path=t_path) |
| result = translator.predict(args.text) |
|
|
| print("\n" + "=" * 50) |
| print(f"Source ({args.direction.upper()}) : {args.text}") |
| print(f"Translation: {result}") |
| print("=" * 50 + "\n") |
|
|
