Update app.py

app.py

@@ -3,14 +3,14 @@ import torch
 from datasets import load_dataset, concatenate_datasets
 from transformers import AutoModelForCausalLM, AutoTokenizer, Trainer, TrainingArguments
 
-# Predefined datasets
-dataset_names = [
-    'imdb',
-    'ag_news',
-    'squad',
-    'cnn_dailymail',
-    'wiki40b'
-]
+# Predefined datasets with their configurations
+dataset_names = {
+    'imdb': None,
+    'ag_news': None,
+    'squad': None,
+    'cnn_dailymail': '1.0.0',  # Specify configuration for cnn_dailymail
+    'wiki40b': None
+}
 
 # Global variables for model and tokenizer
 model = None
@@ -18,7 +18,9 @@ tokenizer = None
 
 # Function to load and prepare datasets
 def load_and_prepare_datasets():
-    datasets = [load_dataset(name) for name in dataset_names]
+    datasets = []
+    for name, config in dataset_names.items():
+        datasets.append(load_dataset(name, config))
     
     # Concatenate train datasets only for training
     train_dataset = concatenate_datasets([ds['train'] for ds in datasets if 'train' in ds])
@@ -109,4 +111,4 @@ with gr.Blocks() as demo:
     generate_button.click(generate_text, inputs=prompt_input, outputs=generated_output)
 
 # Launch the app with share=True to create a public link
-demo.launch(share=True)
+demo.launch(share=True)