| import argparse |
|
|
| import pandas as pd |
| import torch |
|
|
| from src.utils.helpers import TRANSFORMER_MODEL, TOKENIZER_MODEL, JSON_DATA, CSV_DATA |
| from src.models.transformer.seq2seq import BaselineSeq2SeqTransformer |
| from src.models.transformer.tokenizer import ( |
| TranslationDataset, |
| PAD_IDX, |
| SOS_IDX, |
| EOS_IDX, |
| ) |
| from src.models.transformer.helpers import generate_square_subsequent_mask |
|
|
| from scripts.preprocess import preprocess_text |
|
|
| SAMPLE_SIZE = 1500 |
| MAX_SEQ_LEN = 128 |
| BEAM_SIZE = 3 |
| MAX_DECODING_LEN = 40 |
| ALPHA = 0.7 |
|
|
|
|
| def configure_device() -> torch.device: |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
| print(f"Using device: {device}") |
|
|
| if torch.cuda.is_available(): |
| torch.cuda.empty_cache() |
| torch.cuda.reset_peak_memory_stats(device) |
| initial_vram = torch.cuda.memory_allocated(device) / (1024**2) |
| print(f"Initial GPU Memory Usage: {initial_vram:.2f} MB") |
|
|
| return device |
|
|
|
|
| def beam_search_decode( |
| model: torch.nn.Module, |
| src_tensor: torch.Tensor, |
| device: torch.device, |
| beam_size: int = BEAM_SIZE, |
| max_len: int = MAX_DECODING_LEN, |
| alpha: float = ALPHA, |
| ): |
| 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 |
|
|
|
|
| def translate_sentence( |
| sentence: str, model: torch.nn.Module, sp, device: torch.device |
| ) -> str: |
| cleaned_sentence = preprocess_text(sentence.lower()) |
| src_tokens = [SOS_IDX] + sp.encode(str(cleaned_sentence), out_type=int) + [EOS_IDX] |
| src_tensor = torch.tensor(src_tokens, dtype=torch.long).unsqueeze(1) |
|
|
| predicted_ids = beam_search_decode(model, src_tensor, device, beam_size=BEAM_SIZE) |
| clean_ids = [idx for idx in predicted_ids if idx not in [SOS_IDX, EOS_IDX, PAD_IDX]] |
| return sp.decode(clean_ids) |
|
|
|
|
| def run_transformer_batch_inference(num_samples: int = SAMPLE_SIZE): |
| device = configure_device() |
|
|
| print("\nLoading tokenizer and dataset...") |
| train_dataset = TranslationDataset( |
| str(JSON_DATA), |
| spm_model_path=str(TOKENIZER_MODEL), |
| max_seq_len=MAX_SEQ_LEN, |
| ) |
| sp = train_dataset.sp |
| vocab_size = sp.get_piece_size() |
|
|
| 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(device) |
|
|
| model.load_state_dict(torch.load(TRANSFORMER_MODEL, map_location=device)) |
| model.eval() |
| print("Model loaded successfully!") |
|
|
| print("\nLoading test dataset...") |
| df = pd.read_csv(CSV_DATA) |
| if "split" in df.columns: |
| df = df[df["split"] == "test"] |
|
|
| sample_df = df.sample(n=min(num_samples, len(df)), random_state=42).reset_index( |
| drop=True |
| ) |
| print(f"Evaluation samples: {len(sample_df)}") |
|
|
| print("\n" + "=" * 60) |
| print(f"FINAL INFERENCE PIPELINE ({len(sample_df)} SAMPLES)") |
| print("=" * 60) |
|
|
| for idx, row in sample_df.iterrows(): |
| source_sentence = str(row.get("source_text", "")) |
| reference_sentence = str(row.get("target_text", "")) |
|
|
| if not source_sentence: |
| continue |
|
|
| translation = translate_sentence(source_sentence, model, sp, device) |
|
|
| print(f"\nSample {idx + 1}") |
| print(f"Source : {source_sentence}") |
| print(f"Reference : {reference_sentence}") |
| print(f"Prediction : {translation}") |
| print("-" * 60) |
|
|
| print("\nInference pipeline completed successfully!") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser( |
| description="Run Transformer inference on the test split" |
| ) |
| parser.add_argument( |
| "--samples", |
| type=int, |
| default=SAMPLE_SIZE, |
| help="Number of random samples to translate from the test set", |
| ) |
| args = parser.parse_args() |
|
|
| run_transformer_batch_inference(num_samples=args.samples) |
|
|
