| import argparse |
| import time |
| from pathlib import Path |
|
|
| import pandas as pd |
| import torch |
| from sacrebleu.metrics import CHRF |
| from nltk.translate.bleu_score import corpus_bleu, sentence_bleu, SmoothingFunction |
|
|
| 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 |
|
|
| 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: |
| src_tokens = [SOS_IDX] + sp.encode(str(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_evaluation(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() |
|
|
| print("\nBuilding transformer model...") |
| 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) |
|
|
| print("\nLoading model checkpoint...") |
| 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'] |
|
|
| test_df = df.sample(n=min(num_samples, len(df)), random_state=42).reset_index(drop=True) |
| print(f"Evaluation samples: {len(test_df)}") |
|
|
| references = [] |
| predictions = [] |
| results = [] |
| total_inference_time = 0.0 |
| smoothie = SmoothingFunction().method4 |
| chrf_metric = CHRF() |
|
|
| evaluation_start = time.time() |
| print(f"\nStarting evaluation on {len(test_df)} samples...\n") |
|
|
| with torch.no_grad(): |
| for idx, row in test_df.iterrows(): |
| source_sentence = str(row["source_text"]) |
| reference_sentence = str(row["target_text"]) |
|
|
| start_time = time.time() |
| prediction = translate_sentence(source_sentence, model, sp, device) |
| inference_time = time.time() - start_time |
| total_inference_time += inference_time |
|
|
| references.append([reference_sentence.split()]) |
| predictions.append(prediction.split()) |
|
|
| sentence_bleu1 = sentence_bleu( |
| [reference_sentence.split()], |
| prediction.split(), |
| weights=(1, 0, 0, 0), |
| smoothing_function=smoothie, |
| ) |
| sentence_bleu2 = sentence_bleu( |
| [reference_sentence.split()], |
| prediction.split(), |
| weights=(0.5, 0.5, 0, 0), |
| smoothing_function=smoothie, |
| ) |
| sentence_bleu3 = sentence_bleu( |
| [reference_sentence.split()], |
| prediction.split(), |
| weights=(1/3, 1/3, 1/3, 0), |
| smoothing_function=smoothie, |
| ) |
| sentence_bleu4 = sentence_bleu( |
| [reference_sentence.split()], |
| prediction.split(), |
| weights=(0.25, 0.25, 0.25, 0.25), |
| smoothing_function=smoothie, |
| ) |
| sentence_chrf = chrf_metric.sentence_score(prediction, [reference_sentence]).score |
|
|
| results.append({ |
| "source": source_sentence, |
| "reference": reference_sentence, |
| "prediction": prediction, |
| "bleu1": sentence_bleu1, |
| "bleu2": sentence_bleu2, |
| "bleu3": sentence_bleu3, |
| "bleu4": sentence_bleu4, |
| "chrf": sentence_chrf, |
| "inference_time_sec": inference_time, |
| }) |
|
|
| if (idx + 1) % 10 == 0: |
| print(f"Processed {idx + 1}/{len(test_df)} samples") |
|
|
| total_evaluation_time = time.time() - evaluation_start |
| results_df = pd.DataFrame(results) |
|
|
| bleu1 = corpus_bleu( |
| references, |
| predictions, |
| weights=(1, 0, 0, 0), |
| smoothing_function=smoothie, |
| ) |
| bleu2 = corpus_bleu( |
| references, |
| predictions, |
| weights=(0.5, 0.5, 0, 0), |
| smoothing_function=smoothie, |
| ) |
| bleu3 = corpus_bleu( |
| references, |
| predictions, |
| weights=(1/3, 1/3, 1/3, 0), |
| smoothing_function=smoothie, |
| ) |
| bleu4 = corpus_bleu( |
| references, |
| predictions, |
| weights=(0.25, 0.25, 0.25, 0.25), |
| smoothing_function=smoothie, |
| ) |
|
|
| pred_texts = [" ".join(pred) for pred in predictions] |
| ref_texts = [" ".join(ref[0]) for ref in references] |
| chrf_score = chrf_metric.corpus_score(pred_texts, [ref_texts]) |
|
|
| avg_inference_time = total_inference_time / len(test_df) |
| peak_vram = torch.cuda.max_memory_allocated(device) / (1024 ** 2) if torch.cuda.is_available() else 0.0 |
|
|
| print("\n================================================") |
| print("TRANSFORMER MACHINE TRANSLATION EVALUATION") |
| print("================================================") |
| print(f"Evaluated Samples : {len(results_df)}") |
| print(f"BLEU-1 Score : {bleu1:.4f}") |
| print(f"BLEU-2 Score : {bleu2:.4f}") |
| print(f"BLEU-3 Score : {bleu3:.4f}") |
| print(f"BLEU-4 Score : {bleu4:.4f}") |
| print(f"chrF (F-score) : {chrf_score.score:.4f}") |
| print(f"Average Inference Time : {avg_inference_time:.4f} sec/sample") |
| print(f"Total Evaluation Time : {total_evaluation_time:.2f} sec") |
| print(f"Peak VRAM Usage : {peak_vram:.2f} MB") |
|
|
| results_dir = Path(__file__).parent.parent / "results" |
| results_dir.mkdir(parents=True, exist_ok=True) |
|
|
| results_output_path = results_dir / "transformer_mt_evaluation_results_BPE.csv" |
| results_df.to_csv(results_output_path, index=False) |
| print(f"\nDetailed results saved to: {results_output_path}") |
|
|
| summary_metrics = { |
| "BLEU-1": bleu1, |
| "BLEU-2": bleu2, |
| "BLEU-3": bleu3, |
| "BLEU-4": bleu4, |
| "chrF": chrf_score.score, |
| "Average_Inference_Time_sec": avg_inference_time, |
| "Total_Evaluation_Time_sec": total_evaluation_time, |
| "Peak_VRAM_MB": peak_vram, |
| } |
|
|
| summary_df = pd.DataFrame([summary_metrics]) |
| summary_output_path = results_dir / "transformer_evaluation_summary.csv" |
| summary_df.to_csv(summary_output_path, index=False) |
| print(f"Summary metrics saved to: {summary_output_path}") |
| print("\nEvaluation completed successfully!") |
|
|
|
|
| if __name__ == "__main__": |
| parser = argparse.ArgumentParser(description="Run Transformer evaluation 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_evaluation(num_samples=args.samples) |
|
|
