Update main.py

main.py

@@ -13,7 +13,7 @@ from transformers import (AutoTokenizer, BitsAndBytesConfig, MBart50TokenizerFas
                           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 = {
@@ -50,28 +50,6 @@ def summarize_text_mbart50(texts, model, tokenizer):
     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=60,
-        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}")
@@ -99,7 +77,6 @@ def experiments(model_name, experiment_type, num_examples, finetune_type):
     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":
@@ -126,13 +103,10 @@ def run_zero_shot(model_name, model, tokenizer, summarize_text, test, test_fr, t
         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}")
+            print(f"Processing batch {i//batch_size + 1}")
             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)
@@ -166,13 +140,10 @@ def run_1_shot(model_name, model, tokenizer, summarize_text, train, train_fr, tr
         # 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}")
+            print(f"Processing batch {i//batch_size + 1}")
             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)
@@ -208,12 +179,9 @@ def run_2_shot(model_name, model, tokenizer, summarize_text, train, train_fr, tr
         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}")
+            print(f"Processing batch {i//batch_size + 1}")
             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)