gera-richarte commited on
Commit
d91013d
1 Parent(s): bcd3e1f

bugfix on load_parquet()

Browse files
Files changed (1) hide show
  1. earthview.py +3 -2
earthview.py CHANGED
@@ -60,10 +60,11 @@ def load_dataset(subset, dataset="satellogic/EarthView", split="train", shards =
60
 
61
  def load_parquet(subset_or_filename, batch_size=100):
62
  if subset_or_filename in get_subsets():
63
- pqfile = ParquetFile(f"dataset/{subset_or_filename}/sample.parquet")
64
  else:
65
- pqfile = subset_or_filename
66
 
 
67
  batch = pqfile.iter_batches(batch_size=batch_size)
68
  return Dataset(pa_Table.from_batches(batch))
69
 
 
60
 
61
  def load_parquet(subset_or_filename, batch_size=100):
62
  if subset_or_filename in get_subsets():
63
+ filename = f"dataset/{subset_or_filename}/sample.parquet"
64
  else:
65
+ filename = subset_or_filename
66
 
67
+ pqfile = ParquetFile(filename)
68
  batch = pqfile.iter_batches(batch_size=batch_size)
69
  return Dataset(pa_Table.from_batches(batch))
70