Added support for ag_news dataset

README.md

@@ -19,7 +19,7 @@ Tensorflow >= 1.4.0
 
 3. Run `python train.py 20ng`
 
-4. Change `20ng` in above 3 command lines to `R8`, `R52`, `ohsumed` and `mr` when producing results for other datasets.
+4. Change `20ng` in above 3 command lines to `R8`, `R52`, `ohsumed`, `ag_news` and `mr` when producing results for other datasets.
 
 ## Example input data
 

build_graph.py

@@ -15,7 +15,7 @@
 if len(sys.argv) != 2:
 	sys.exit("Use: python build_graph.py <dataset>")
 
-datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr']
+datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr', 'ag_news']
 # build corpus
 dataset = sys.argv[1]
 

remove_words.py

@@ -8,7 +8,7 @@
 if len(sys.argv) != 2:
 	sys.exit("Use: python remove_words.py <dataset>")
 
-datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr']
+datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr', 'ag_news']
 dataset = sys.argv[1]
 
 if dataset not in datasets:

train.py

@@ -14,7 +14,7 @@
 if len(sys.argv) != 2:
 	sys.exit("Use: python train.py <dataset>")
 
-datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr']
+datasets = ['20ng', 'R8', 'R52', 'ohsumed', 'mr', 'ag_news']
 dataset = sys.argv[1]
 
 if dataset not in datasets: