修改测试文件

tensnsorflow/record.py

@@ -54,18 +54,15 @@ def gen_tfrecords(in_file):
 
 def main(_):
     client = Client("http://nvwa01:50070")
-    path = "/recommend/va"
-    files = client.walk(path)
-    print(files[3])
-    print(type(files[3]))
-    # for root, dir, files in client.walk(path):
-    #     for file in files:
-    #         print(file)
+    file_list = []
+    for root, dir, files in client.walk(FLAGS.input_dir):
+        for file in files:
+            if file[-5:] == "*.avro":
+                file_list.append(FLAGS.input_dir+file)
 
     if not os.path.exists(FLAGS.output_dir):
         os.mkdir(FLAGS.output_dir)
-    file_list = glob.glob(os.path.join(FLAGS.input_dir, "*.avro"))
-    print(FLAGS.input_dir)
+
     print(file_list)
     print("total files: %d" % len(file_list))