Skip to content

Commit

Permalink
add progress bar to sampler
Browse files Browse the repository at this point in the history
  • Loading branch information
kjappelbaum committed Aug 15, 2024
1 parent d8a6110 commit c0ddf9f
Showing 1 changed file with 6 additions and 2 deletions.
8 changes: 6 additions & 2 deletions src/chemnlp/data/sampler.py
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@
import yaml
import json
from loguru import logger
from tqdm import tqdm


class TemplateSampler:
Expand Down Expand Up @@ -836,8 +837,11 @@ def export(self, output_dir: str, template: str) -> pd.DataFrame:
)
for split in self.df["split"].unique():
df_split = self.df[self.df["split"] == split]
samples = [self.sample(row, template) for _, row in df_split.iterrows()]

samples = []
for _, row in tqdm(df_split.iterrows(), total=len(df_split)):
sample_dict = row.to_dict()
sample = self._fill_template(template, sample_dict)
samples.append(sample)
df_out = pd.DataFrame(samples)

# if self.benchmarking_templates:
Expand Down

0 comments on commit c0ddf9f

Please sign in to comment.