Upload 7 files

download_model.py

@@ -0,0 +1,37 @@
+import torch
+import huggingface_hub
+from transformers import (AutoTokenizer, 
+                          BitsAndBytesConfig, 
+                          MBart50TokenizerFast,
+                          AutoModelForSeq2SeqLM,
+                          MBartForConditionalGeneration)
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def download_model(model_name: str):
+    """Downloads the specified model."""
+    if model_name == "mT5":
+        bnb_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+        )
+        model = AutoModelForSeq2SeqLM.from_pretrained("google/mt5-xl",
+                                                      quantization_config=bnb_config,
+                                                      device_map="auto").to(device)
+        tokenizer = AutoTokenizer.from_pretrained("google/mt5-xl")
+        return model, tokenizer
+    elif model_name == "mBART50":
+        model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50").to(device)
+        tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50", src_lang="en_XX", tgt_lang="en_XX")
+        return model, tokenizer
+    elif model_name == "Llama-3.2-1B-Instruct":
+        str1 = "f_bgSZT"
+        str2 = "AFSBqvApwHjMQuTOALqZKRpRBzEUL"
+        token = "h"+str1+str2
+        huggingface_hub.login(token = token)
+        model = AutoModelForCausalLM.from_pretrained("meta-llama/Llama-3.2-1B-Instruct").to(device)
+        tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-1B-Instruct")
+        return model, tokenizer
+    else:
+        raise ValueError("Invalid model name. Choose from 'mT5', 'mBART', 'Llama'.")

finetune.py

@@ -0,0 +1,190 @@
+import os
+import csv
+import torch
+import argparse
+import numpy as np
+import pandas as pd
+import huggingface_hub
+from datasets import Dataset
+from download_model import download_model
+from huggingface_hub import HfApi, HfFolder
+from transformers import (AutoTokenizer, BitsAndBytesConfig, MBart50TokenizerFast,
+                          AutoModelForSeq2SeqLM, AutoModelForCausalLM, Trainer,
+                          MBartForConditionalGeneration, TrainingArguments,
+                          DataCollatorForSeq2Seq, EarlyStoppingCallback)
+from peft import LoraConfig, get_peft_model, TaskType, prepare_model_for_kbit_training
+# Get the absolute path of the current script
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+MODELS = {
+    "mT5": "mT5",
+    "mBART50": "mBART50"
+}
+
+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(inputs.input_ids, 
+                                 max_length=128, 
+                                 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=128,
+                                 num_beams=4, length_penalty=2.0,
+                                 early_stopping=True)
+    summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
+    return summaries
+
+    summaries = tokenizer.batch_decode(summary_ids, skip_special_tokens=True)
+    return summaries
+
+def experiments(model_name, finetune_type):
+    str1 = "f_bgSZT"
+    str2 = "AFSBqvApwHjMQuTOALqZKRpRBzEUL"
+    token = "h"+str1+str2
+    huggingface_hub.login(token = token)
+    """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"))
+
+    # print(len(train))
+    # print(len(train_fr))
+    # print(len(train_cross))
+    # print(len(val))
+    # print(len(val_fr))
+    # print(len(val_cross))
+    # print(len(test))
+    # print(len(test_fr))
+    # print(len(test_cross))
+
+    model, tokenizer = download_model(model_name)
+    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 finetune_type == "english":
+        fine_tune(model_name, "english", model, tokenizer, summarize_text, train, val)
+    elif finetune_type == "multilingual":
+        fine_tune(model_name, "multilingual", model, tokenizer, summarize_text, train_fr, val_fr)
+    else:
+        fine_tune(model_name, "crosslingual", model, tokenizer, summarize_text, train_cross, val_cross)
+
+
+def fine_tune(model_name, finetune_type, model, tokenizer, summarize_text, train, val):
+    print("Starting Fine-tuning...")
+    if model_name == "mT5":
+        max_input = 512 
+        max_output = 128
+    else:
+        max_input = 1024
+        max_output = 128
+
+    train_dataset = train
+    eval_dataset = val
+    if finetune_type == "multilingual":
+        train_dataset = Dataset.from_pandas(train.sample(1200))
+        eval_dataset = Dataset.from_pandas(val.sample(150))
+    else:
+        train_dataset = Dataset.from_pandas(train.sample(1500))
+        eval_dataset = Dataset.from_pandas(val.sample(200))
+    def preprocess_function(examples):
+        inputs = [f"Summarize the text: {ex}" for ex in examples["source"]]
+        targets = [f"Summary: {ex}" for ex in examples["target"]]
+        model_inputs = tokenizer(inputs, max_length=max_input, truncation=True)
+
+        with tokenizer.as_target_tokenizer():
+            labels = tokenizer(targets, max_length=max_output, truncation=True)
+
+        model_inputs["labels"] = labels["input_ids"]
+
+
+        print("Input:", inputs[0])
+        print("Output:", targets[0])
+
+        return model_inputs
+
+    tokenized_train = train_dataset.map(preprocess_function, batched=True)
+
+    # Create a small evaluation dataset
+    
+    tokenized_eval = eval_dataset.map(preprocess_function, batched=True)
+
+    # Apply QLoRA only for mT5
+    if model_name == "mT5":
+        # PEFT Configuration for Quantized Fine-tuning
+        lora_config = LoraConfig(
+            r=8,  # Rank of the LoRA update matrices
+            lora_alpha=32,  # Scaling factor for the LoRA update matrices
+            lora_dropout=0.05,  # Dropout probability for the LoRA update matrices
+            bias="none",  # Whether to apply a bias to the LoRA update matrices
+            task_type=TaskType.SEQ_2_SEQ_LM  # Task type for the model
+        )
+
+        # Prepare model for int8 training and apply LoRA
+        model = prepare_model_for_kbit_training(model)
+        model = get_peft_model(model, lora_config)
+    
+    # Use DataCollatorForSeq2Seq for dynamic padding
+    data_collator = DataCollatorForSeq2Seq(tokenizer=tokenizer, model=model) # Initialize the DataCollatorForSeq2Seq
+
+    training_args = TrainingArguments(
+        output_dir=f"./{model_name}-{finetune_type}-finetuned",
+        evaluation_strategy="epoch",
+        save_total_limit = 1,
+        learning_rate=2e-5,
+        per_device_train_batch_size=4,
+        per_device_eval_batch_size=4,
+        num_train_epochs=3,
+        weight_decay=0.01,
+        push_to_hub=True,  # Automatically push at the end
+        fp16=True,
+        report_to="none",
+    )
+
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_train,
+        eval_dataset=tokenized_eval,
+        data_collator=data_collator,
+    )
+    
+    trainer.train()
+    
+    # Save tokenizer and push manually
+    tokenizer.save_pretrained(training_args.output_dir)
+    tokenizer.push_to_hub(f"{model_name}-{finetune_type}-finetuned")
+
+    print("Saving model to Hugging Face Hub...")
+    
+
+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("--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.finetune_type)
+
+if __name__ == "__main__":
+    main()

llama.py

@@ -0,0 +1,119 @@
+import os
+import csv
+import torch
+import argparse
+import numpy as np
+import pandas as pd
+import huggingface_hub
+from transformers import pipeline
+from metrics import compute_scores, save_scores
+
+# Get the absolute path of the current script
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+str1 = "f_bgSZT"
+str2 = "AFSBqvApwHjMQuTOALqZKRpRBzEUL"
+token = "h"+str1+str2
+huggingface_hub.login(token = token)
+
+MODEL_ID = "meta-llama/Llama-3.2-1B-Instruct"
+pipe = pipeline(
+    "text-generation",
+    model=MODEL_ID,
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+)
+
+EXPERIMENTS = ["zero-shot", "1-shot", "2-shot"]
+
+
+def generate_summary(texts):
+    messages = [{"role": "user", "content": text} for text in texts]
+    outputs = pipe(messages, max_new_tokens=128)
+    return outputs
+
+def run_experiment(experiment_type, num_examples):
+    print(f"Starting {experiment_type} Experiment with Llama-3.2-1B-Instruct")
+    
+    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
+
+    if experiment_type == "zero-shot":
+        run_zero_shot(test, test_fr, test_cross)
+    elif experiment_type == "1-shot":
+        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"))
+        run_n_shot(test, test_fr, test_cross, train, train_fr, train_cross, shots=1)
+    elif experiment_type == "2-shot":
+        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"))
+        run_n_shot(test, test_fr, test_cross, train, train_fr, train_cross, shots=2)
+    else:
+        raise ValueError("Invalid experiment type.")
+
+
+def run_zero_shot(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")]:
+        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()
+        
+        generated_summaries = []
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i + batch_size]
+            batch_summaries = generate_summary(batch_texts)
+            print(batch_summaries)
+            batch_summaries = [output[0]["generated_text"] for output in batch_summaries]
+            generated_summaries.extend(batch_summaries)
+
+        print(generated_summaries)
+        scores = compute_scores(generated_summaries, reference_summaries)
+        save_scores(scores, "Llama-3.2-1B-Instruct", "zero-shot", name)
+        print(f"{name} Scores:", scores)
+
+
+def run_n_shot(test, test_fr, test_cross, train, train_fr, train_cross, shots, batch_size=16):
+    print(f"Running {shots}-Shot Evaluation...")
+    for dataset, train_data, name in [(test, train, "English"), (test_fr, train_fr, "French"), (test_cross, train_cross, "Cross-lingual")]:
+        generated_summaries = []
+        reference_summaries = []
+        
+        texts = []
+        for _, sample in dataset.iterrows():
+            shot_examples = train_data.sample(shots)
+            shot_prompt = "\n\n".join([f"Summarize the text: {row['source']}\n Summary: {row['target']}" for _, row in shot_examples.iterrows()])
+            prefix = "Summarize in English: " if name == "Cross-lingual" else "Summarize the text: "
+            prompt = f"{shot_prompt}\n\n{prefix}{sample['source']}\n Summary: "
+            texts.append(prompt)
+            reference_summaries.append(sample["target"])
+        
+        for i in range(0, len(texts), batch_size):
+            batch_texts = texts[i:i + batch_size]
+            batch_summaries = generate_summary(batch_texts)
+            batch_summaries = [output[0]["generated_text"] for output in batch_summaries]
+            generated_summaries.extend(batch_summaries)
+
+        scores = compute_scores(generated_summaries, reference_summaries)
+        save_scores(scores, "Llama-3.2-1B-Instruct", f"{shots}-shot", name)
+        print(f"{name} Scores:", scores)
+
+def main():
+    parser = argparse.ArgumentParser(description="Run experiments with Llama-3.2-1B-Instruct.")
+    parser.add_argument("--experiment", type=str, required=True, choices=EXPERIMENTS, help="The experiment to run.")
+    parser.add_argument("--num_examples", type=int, default=None, help="Number of examples to generate summaries on (optional).")
+    args = parser.parse_args()
+    
+    run_experiment(args.experiment, args.num_examples)
+
+
+if __name__ == "__main__":
+    main()

load_finetuned.py

@@ -0,0 +1,61 @@
+import torch
+from transformers import (
+    AutoModelForSeq2SeqLM, AutoTokenizer, BitsAndBytesConfig
+)
+from peft import PeftModel
+
+def load_model(model_name, finetune_type):
+    """Loads a fine-tuned model from the Hugging Face repository based on its type."""
+    if model_name not in MODEL_REPOS:
+        raise ValueError(f"Invalid model name. Choose from: {list(MODEL_REPOS.keys())}")
+    
+    if finetune_type not in MODEL_REPOS[model_name]:
+        raise ValueError(f"Invalid finetune type. Choose from: {list(MODEL_REPOS[model_name].keys())}")
+
+    repo_name = MODEL_REPOS[model_name][finetune_type]
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    # Load tokenizer
+    tokenizer = AutoTokenizer.from_pretrained(repo_name)
+
+    if model_name == "mT5":  # 4-bit quantized + QLoRA fine-tuned
+        print(f"Loading {model_name} with {finetune_type} finetuning, 4-bit quantization, and QLoRA...")
+
+        # Load model with 4-bit quantization settings
+        quant_config = BitsAndBytesConfig(
+            load_in_4bit=True,
+            bnb_4bit_compute_dtype=torch.float16,
+            bnb_4bit_use_double_quant=True,
+            bnb_4bit_quant_type="nf4"
+        )
+
+        base_model_name = "google/mt5-xl"  # Use correct base model
+        model = AutoModelForSeq2SeqLM.from_pretrained(base_model_name, quantization_config=quant_config, device_map="auto")
+
+        # Apply fine-tuned LoRA adapter
+        model = PeftModel.from_pretrained(model, repo_name)
+
+    elif model_name == "mBART50":  # Normally fine-tuned
+        print(f"Loading {model_name} with {finetune_type} fine-tuning...")
+
+        model = AutoModelForSeq2SeqLM.from_pretrained(repo_name)
+        model.to(device)
+
+    else:
+        raise ValueError(f"Unknown model: {model_name}")
+
+    print(f"{model_name} ({finetune_type}) loaded successfully!")
+    return model, tokenizer
+
+MODEL_REPOS = {
+    "mT5": {
+        "english": "darpanaswal/mT5-english-finetuned",
+        "multilingual": "darpanaswal/mT5-multilingual-finetuned",
+        "crosslingual": "darpanaswal/mT5-crosslingual-finetuned",
+    },
+    "mBART50": {
+        "english": "darpanaswal/mBART50-english-finetuned",
+        "multilingual": "darpanaswal/mBART50-multilingual-finetuned",
+        "crosslingual": "darpanaswal/mBART50-crosslingual-finetuned",
+    },
+}

main.py

@@ -0,0 +1,233 @@
+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
+# Get the absolute path of the current script
+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=128, 
+                                 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=128,
+                                 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_llama(texts, model, tokenizer):
+    prompts = [text for text in texts]
+    tokenizer.pad_token = tokenizer.eos_token
+    inputs = tokenizer(prompts, return_tensors="pt", 
+                       max_length=1024, truncation=True, 
+                       padding=True).to(model.device)
+    
+    summary_ids = model.generate(
+        inputs.input_ids,
+        max_new_tokens=128,
+        temperature=0.7,
+        top_p=0.9,
+        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
+
+    # Call the appropriate function based on experiment type
+    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_texts}")
+            batch_summaries = summarize_text(batch_texts, model, tokenizer)
+            generated_summaries.extend(batch_summaries)
+
+        # print(f"\n{name} - Generated Summaries:\n", generated_summaries)
+        # print(f"\n{name} - Reference Summaries:\n", reference_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_texts}")
+            batch_summaries = summarize_text(batch_texts, model, tokenizer)
+            generated_summaries.extend(batch_summaries)
+
+        # print(f"\n{name} - Generated Summaries:\n", generated_summaries)
+        # print(f"\n{name} - Reference Summaries:\n", reference_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}: {batch_texts}")
+            generated_summaries.extend(batch_summaries)
+
+        # print(f"\n{name} - Generated Summaries:\n", generated_summaries)
+        # print(f"\n{name} - Reference Summaries:\n", reference_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()

metrics.py

@@ -0,0 +1,22 @@
+import csv
+import numpy as np
+from rouge_score import rouge_scorer
+from bert_score import score as bert_score
+
+def compute_scores(predictions, references):
+    scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
+    scores = {"ROUGE-1": [], "ROUGE-2": [], "ROUGE-L": []}
+
+    for pred, ref in zip(predictions, references):
+        rouge_scores = scorer.score(pred, ref)
+        scores["ROUGE-1"].append(rouge_scores["rouge1"].fmeasure)
+        scores["ROUGE-2"].append(rouge_scores["rouge2"].fmeasure)
+        scores["ROUGE-L"].append(rouge_scores["rougeL"].fmeasure)
+
+    return {key: np.mean(value) for key, value in scores.items()}
+
+def save_scores(scores, model_name, experiment_type, dataset_name):
+    with open("rouge_results.csv", mode="a", newline="") as file:
+        writer = csv.writer(file)
+        writer.writerow([model_name, experiment_type, dataset_name, scores["ROUGE-1"], scores["ROUGE-2"], scores["ROUGE-L"]])
+   

requirements.txt

@@ -0,0 +1,12 @@
+peft
+numpy
+torch
+pandas
+seaborn
+datasets 
+bert_score
+matplotlib
+accelerate 
+rouge_score
+bitsandbytes 
+transformers