| 14 | |
| 15 | |
| 16 | def extractVectors(dfs, ds_path, dataFolder, augment, sample, table_order, run_id, singleCol=False): |
| 17 | if singleCol: |
| 18 | model_path = "/home/benchmark/starmie-main/" \ |
| 19 | "model_%s_%s_%s_%dsingleCol.pt" % (augment, sample, table_order, run_id) |
| 20 | else: |
| 21 | model_path = "/home/benchmark/starmie-main/" \ |
| 22 | "model_%s_%s_%s_%d.pt" % (augment, sample, table_order, run_id) |
| 23 | print(f"model_path: {model_path}") |
| 24 | ckpt = torch.load(model_path, map_location=torch.device('cuda')) |
| 25 | # load_checkpoint from sdd/pretain |
| 26 | |
| 27 | model, trainset = load_checkpoint(ckpt, ds_path) |
| 28 | |
| 29 | return inference_on_tables(dfs, model, trainset, batch_size=528) |
| 30 | |
| 31 | |
| 32 | def get_df(dataFolder): |