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': 'en'  # Specify language for wiki40b
+    'wiki40b': 'ru'  # Specify language for wiki40b
 }
 
 # Global variables for model and tokenizer
@@ -22,11 +22,21 @@ def load_and_prepare_datasets():
     for name, config in dataset_names.items():
         datasets.append(load_dataset(name, config))
     
+    # Extract only the 'text' field 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 'test' in ds:
+            eval_datasets.append(ds['test'].map(lambda x: {'text': x['text']}))
+    
     # Concatenate train datasets only for training
-    train_dataset = concatenate_datasets([ds['train'] for ds in datasets if 'train' in ds])
+    train_dataset = concatenate_datasets(train_datasets)
     
-    # Use only a subset for evaluation if needed
-    eval_dataset = concatenate_datasets([ds['test'] for ds in datasets if 'test' in ds])
+    # Concatenate eval datasets only for evaluation
+    eval_dataset = concatenate_datasets(eval_datasets)
     
     return train_dataset, eval_dataset