Merge pull request #47 from OpenManus/test-pipeline

Test pipeline add first version

Author: Kunlun-Zhu

data/generate_train_webshop.py

@@ -78,9 +78,9 @@ def flatten_data(items_data):
     parser = argparse.ArgumentParser()
     parser.add_argument('--input_file', default="openmanus_rl/agentgym/agentenv-webshop/webshop/data/items_human_ins.json", 
                         help="Path to items_human_ins.json")
-    parser.add_argument('--output_dir', required=True, help="Output directory for processed parquet")
+    parser.add_argument('--output_dir', required=False, default="data/webshop", help="Output directory for processed parquet")
     parser.add_argument('--split', type=str, default="train")
-    parser.add_argument('--train_ratio', type=float, default=0.95, 
+    parser.add_argument('--train_ratio', type=float, default=0.90, 
                         help="Ratio of data to use for training (rest for val/test)")
     parser.add_argument('--val_ratio', type=float, default=0.1,
                         help="Ratio of data to use for validation")
@@ -99,6 +99,11 @@ def flatten_data(items_data):
     if args.split == "all":
         # Process all data with the same split label
         dataset = dataset.map(function=make_map_fn(args.split), with_indices=True)
+        # Save the entire dataset
+        output_path = os.path.join(args.output_dir, f"{args.split}.parquet")
+        dataset.to_parquet(output_path)
+        print(f"Processed {len(dataset)} examples")
+        print(f"Data saved to {output_path}")
     else:
         # Split the dataset
         splits = dataset.train_test_split(
@@ -108,30 +113,43 @@ def flatten_data(items_data):
 
         # Further split the test set into validation and test
         if args.val_ratio > 0:
+            # Calculate the ratio for the validation set from the remaining data
+            remaining_ratio = 1.0 - args.train_ratio
+            val_test_ratio = max(0.5, args.val_ratio / remaining_ratio)  # Ensure ratio is valid
             test_val_split = splits["test"].train_test_split(
-                test_size=(1.0 - (args.val_ratio / (1.0 - args.train_ratio))),
+                test_size=0.5,  # Split remaining data equally between val and test
                 seed=42
             )
             splits = {
                 "train": splits["train"],
                 "validation": test_val_split["train"],
                 "test": test_val_split["test"]
             }
-        
-        # Process only the requested split
-        if args.split in splits:
-            dataset = splits[args.split]
-            dataset = dataset.map(function=make_map_fn(args.split), with_indices=True)
+            
+            # Process and save all splits
+            for split_name, split_dataset in splits.items():
+                processed_dataset = split_dataset.map(function=make_map_fn(split_name), with_indices=True)
+                output_path = os.path.join(args.output_dir, f"{split_name}.parquet")
+                processed_dataset.to_parquet(output_path)
+                print(f"Processed {len(processed_dataset)} examples for {split_name}")
+                print(f"Data saved to {output_path}")
+                
+                # Print sample for the requested split
+                if split_name == args.split:
+                    dataset = processed_dataset
         else:
-            raise ValueError(f"Invalid split: {args.split}. Must be 'train', 'validation', 'test', or 'all'")
+            # If no validation split is requested, just process train and test
+            for split_name, split_dataset in splits.items():
+                processed_dataset = split_dataset.map(function=make_map_fn(split_name), with_indices=True)
+                output_path = os.path.join(args.output_dir, f"{split_name}.parquet")
+                processed_dataset.to_parquet(output_path)
+                print(f"Processed {len(processed_dataset)} examples for {split_name}")
+                print(f"Data saved to {output_path}")
+                
+                # Set the dataset for the requested split
+                if split_name == args.split:
+                    dataset = processed_dataset
 
-    # Create output directory and save dataset
-    os.makedirs(args.output_dir, exist_ok=True)
-    output_path = os.path.join(args.output_dir, f"{args.split}.parquet")
-    dataset.to_parquet(output_path)
-    
-    print(f"Processed {len(dataset)} examples")
-    print(f"Data saved to {output_path}")
-    
-    # Print sample
+    # Print sample from the requested split
+    print(f"\nSample from {args.split} split:")
     pprint(dataset[0])