Update app.py

app.py

@@ -9,7 +9,7 @@ dataset_names = {
     'ag_news': None,
     'squad': None,
     'cnn_dailymail': '1.0.0',  # Specify configuration for cnn_dailymail
-    'wiki40b': 'ru'  # Specify language for wiki40b
+    'wiki40b': 'en'  # Specify language for wiki40b
 }
 
 # Global variables for model and tokenizer
@@ -20,18 +20,33 @@ tokenizer = None
 def load_and_prepare_datasets():
     datasets = []
     for name, config in dataset_names.items():
-        datasets.append(load_dataset(name, config))
-    
-    # Extract only the 'text' field from each dataset for training
+        ds = load_dataset(name, config)
+        datasets.append(ds)
+
+        # Print dataset features for debugging
+        print(f"Dataset: {name}, Features: {ds['train'].features}")
+
+    # Extract only the relevant fields from each dataset for training
     train_datasets = []
     eval_datasets = []
     
     for ds in datasets:
         if 'train' in ds:
-            train_datasets.append(ds['train'].map(lambda x: {'text': x['text']}))
+            if 'text' in ds['train'].features:
+                train_datasets.append(ds['train'].map(lambda x: {'text': x['text']}))
+            elif 'content' in ds['train'].features:  # Example for CNN/DailyMail
+                train_datasets.append(ds['train'].map(lambda x: {'text': x['content']}))
+            else:
+                print(f"Warning: No suitable text field found in {ds['train'].features}")
+        
         if 'test' in ds:
-            eval_datasets.append(ds['test'].map(lambda x: {'text': x['text']}))
-    
+            if 'text' in ds['test'].features:
+                eval_datasets.append(ds['test'].map(lambda x: {'text': x['text']}))
+            elif 'content' in ds['test'].features:  # Example for CNN/DailyMail
+                eval_datasets.append(ds['test'].map(lambda x: {'text': x['content']}))
+            else:
+                print(f"Warning: No suitable text field found in {ds['test'].features}")
+
     # Concatenate train datasets only for training
     train_dataset = concatenate_datasets(train_datasets)