Update app.py

app.py

@@ -37,19 +37,15 @@ def load_or_finetune_pegasus():
         tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-xsum")
         model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
         
-        cnn_dm = load_dataset("cnn_dailymail", "3.0.0", split="train[:5000]")
-        xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True)
+        # Load and normalize datasets
+        cnn_dm = load_dataset("cnn_dailymail", "3.0.0", split="train[:5000]").rename_column("article", "text").rename_column("highlights", "summary")
+        xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True).rename_column("document", "text")
         combined_dataset = concatenate_datasets([cnn_dm, xsum])
         
         def preprocess_function(examples):
-            # Ensure texts and summaries are lists of strings
-            texts = examples["article"] if "article" in examples else examples["document"]
-            summaries = examples["highlights"] if "highlights" in examples else examples["summary"]
-            
-            # Tokenize inputs and targets
-            inputs = tokenizer(texts, max_length=512, truncation=True, padding="max_length", return_tensors="pt")
-            targets = tokenizer(summaries, max_length=400, truncation=True, padding="max_length", return_tensors="pt")
-            
+            # Directly use normalized 'text' and 'summary' fields
+            inputs = tokenizer(examples["text"], max_length=512, truncation=True, padding="max_length", return_tensors="pt")
+            targets = tokenizer(examples["summary"], max_length=400, truncation=True, padding="max_length", return_tensors="pt")
             inputs["labels"] = targets["input_ids"]
             return inputs