import os import csv import torch import argparse import numpy as np import pandas as pd import huggingface_hub from datasets import Dataset from load_finetuned import load_model from metrics import compute_scores, save_scores from transformers import (AutoTokenizer, BitsAndBytesConfig, MBart50TokenizerFast, AutoModelForSeq2SeqLM, AutoModelForCausalLM, Trainer, MBartForConditionalGeneration, TrainingArguments, DataCollatorForSeq2Seq) from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training BASE_DIR = os.path.dirname(os.path.abspath(__file__)) MODELS = { "mT5": "mT5", "mBART50": "mBART50" } LANGUAGE_CODES = { "English": "en_XX", "French": "fr_XX" } EXPERIMENTS = { "mT5": ["zero-shot"], "mBART50": ["zero-shot", "1-shot"] } def summarize_text_mt5(texts, model, tokenizer): inputs = tokenizer(texts, return_tensors="pt", max_length=512, truncation=True, padding=True).to(model.device) summary_ids = model.generate(input_ids = inputs.input_ids, max_length=60, num_beams=4, length_penalty=2.0, early_stopping=True) summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True) return summaries def summarize_text_mbart50(texts, model, tokenizer): inputs = tokenizer(texts, return_tensors="pt", max_length=1024, truncation=True, padding=True).to(model.device) summary_ids = model.generate(inputs.input_ids, max_length=60, num_beams=4, length_penalty=2.0, early_stopping=True) summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True) return summaries def experiments(model_name, experiment_type, num_examples, finetune_type): """Runs an experiment with the given model and dataset.""" print(f"Starting Experiment: on {model_name}") # Construct dataset paths dynamically train = pd.read_csv(os.path.join(BASE_DIR, "datasets/train.csv")) train_fr = pd.read_csv(os.path.join(BASE_DIR, "datasets/train_fr.csv")) train_cross = pd.read_csv(os.path.join(BASE_DIR, "datasets/train_cross.csv")) val = pd.read_csv(os.path.join(BASE_DIR, "datasets/val.csv")) val_fr = pd.read_csv(os.path.join(BASE_DIR, "datasets/val_fr.csv")) val_cross = pd.read_csv(os.path.join(BASE_DIR, "datasets/val_cross.csv")) test = pd.read_csv(os.path.join(BASE_DIR, "datasets/test.csv")) test_fr = pd.read_csv(os.path.join(BASE_DIR, "datasets/test_fr.csv")) test_cross = pd.read_csv(os.path.join(BASE_DIR, "datasets/test_cross.csv")) test = test.sample(num_examples) if num_examples else test test_fr = test_fr.sample(num_examples) if num_examples else test_fr test_cross = test_cross.sample(num_examples) if num_examples else test_cross model, tokenizer = load_model(model_name, finetune_type) print(f"Model {model_name} loaded successfully.") if model_name == "mT5": summarize_text = summarize_text_mt5 elif model_name == "mBART50": summarize_text = summarize_text_mbart50 if experiment_type == "zero-shot": run_zero_shot(model_name, model, tokenizer, summarize_text, test, test_fr, test_cross) elif experiment_type == "1-shot": run_1_shot(model_name, model, tokenizer, summarize_text, train, train_fr, train_cross, test, test_fr, test_cross) elif experiment_type == "2-shot": run_2_shot(model_name, model, tokenizer, summarize_text, train, train_fr, train_cross, test, test_fr, test_cross) else: raise ValueError("Invalid experiment type.") def run_zero_shot(model_name, model, tokenizer, summarize_text, test, test_fr, test_cross, batch_size=16): print("Running Zero-Shot Evaluation...") for dataset, name in [(test, "English"), (test_fr, "French"), (test_cross, "Cross-lingual")]: if model_name == "mBART50": if name == "English": tokenizer.src_lang = "en_XX" else: tokenizer.src_lang = "fr_XX" prefix = "Summarize in English: " if name == "Cross-lingual" else "Summarize the text: " texts = [f"{prefix}{row['source']}\n Summary: " for _, row in dataset.iterrows()] reference_summaries = dataset["target"].tolist() # Process in batches generated_summaries = [] for i in range(0, len(texts), batch_size): batch_texts = texts[i:i + batch_size] print(f"Processing batch {i//batch_size + 1}") batch_summaries = summarize_text(batch_texts, model, tokenizer) generated_summaries.extend(batch_summaries) scores = compute_scores(generated_summaries, reference_summaries) save_scores(scores, model_name, "zero-shot", name) print(f"{name} Scores:", scores) def run_1_shot(model_name, model, tokenizer, summarize_text, train, train_fr, train_cross, test, test_fr, test_cross, batch_size=16): print("Running 1-Shot Evaluation...") for dataset, train_data, name in [(test, train, "English"), (test_fr, train_fr, "French"), (test_cross, train_cross, "Cross-lingual")]: if model_name == "mBART50": if name == "English": tokenizer.src_lang = "en_XX" else: tokenizer.src_lang = "fr_XX" generated_summaries = [] reference_summaries = [] texts = [] for _, sample in dataset.iterrows(): one_shot = train_data.sample(1) source = one_shot["source"].iloc[0] target = one_shot["target"].iloc[0] prefix = "Summarize in English: " if name == "Cross-lingual" else "Summarize the text: " prompt = ( f"{prefix}{source}\n Summary: {target}\n\n" f"{prefix}{sample['source']}\n Summary: " ) texts.append(prompt) reference_summaries.append(sample["target"]) # Process in batches for i in range(0, len(texts), batch_size): batch_texts = texts[i:i + batch_size] print(f"Processing batch {i//batch_size + 1}") batch_summaries = summarize_text(batch_texts, model, tokenizer) generated_summaries.extend(batch_summaries) scores = compute_scores(generated_summaries, reference_summaries) save_scores(scores, model_name, "1-shot", name) print(f"{name} Scores:", scores) def run_2_shot(model_name, model, tokenizer, summarize_text, train, train_fr, train_cross, test, test_fr, test_cross, batch_size=16): print("Running 2-Shot Evaluation...") for dataset, train_data, name in [(test, train, "English"), (test_fr, train_fr, "French"), (test_cross, train_cross, "Cross-lingual")]: if model_name == "mBART50": if name == "English": tokenizer.src_lang = "en_XX" else: tokenizer.src_lang = "fr_XX" generated_summaries = [] reference_summaries = [] texts = [] for _, sample in dataset.iterrows(): two_shots = train_data.sample(2) two_shot1, two_shot2 = two_shots.iloc[0], two_shots.iloc[1] source1, source2 = two_shot1["source"].iloc[0] , two_shot2["source"].iloc[0] target1, target2 = two_shot1["target"].iloc[0] , two_shot2["target"].iloc[0] prefix = "Summarize in English: " if name == "Cross-lingual" else "Summarize the text: " prompt = ( f"{prefix}{two_shot1['source']}\n Summary: {two_shot1['target']}\n\n" f"{prefix}{two_shot2['source']}\n Summary: {two_shot2['target']}\n\n" f"{prefix}{sample['source']}\n Summary: " ) texts.append(prompt) reference_summaries.append(sample["target"]) # Process in batches for i in range(0, len(texts), batch_size): batch_texts = texts[i:i + batch_size] batch_summaries = summarize_text(batch_texts, model, tokenizer) print(f"Processing batch {i//batch_size + 1}") generated_summaries.extend(batch_summaries) scores = compute_scores(generated_summaries, reference_summaries) save_scores(scores, model_name, "2-shot", name) print(f"{name} Scores:", scores) def main(): parser = argparse.ArgumentParser(description="Run experiments with different models.") parser.add_argument("--model", type=str, required=True, choices=MODELS.values(), help="The model to use.") parser.add_argument("--experiment", type=str, required=True, choices=sum(EXPERIMENTS.values(), []), help="The experiment to run.") parser.add_argument("--num_examples", type=int, default=None, help="Number of examples to generate summaries on (optional).") parser.add_argument("--finetune_type", type=str, required=True, choices=["english", "multilingual", "crosslingual"], help="The type of fine-tuning to apply.") args = parser.parse_args() experiments(args.model, args.experiment, args.num_examples, args.finetune_type) if __name__ == "__main__": main()