Update app.py

app.py

@@ -37,15 +37,14 @@ def load_or_finetune_pegasus():
         tokenizer = PegasusTokenizer.from_pretrained("google/pegasus-xsum")
         model = PegasusForConditionalGeneration.from_pretrained("google/pegasus-xsum")
         
-        # 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")
+        cnn_dm = load_dataset("cnn_dailymail", "3.0.0", split="train[:5000]")
+        xsum = load_dataset("xsum", split="train[:5000]", trust_remote_code=True)
         combined_dataset = concatenate_datasets([cnn_dm, xsum])
         
         def preprocess_function(examples):
-            # Ensure texts and summaries are List[str], filter out None or non-string values
-            texts = [str(t) for t in examples["text"] if t is not None]
-            summaries = [str(s) for s in examples["summary"] if s is not None]
+            # 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")