| 1314 | |
| 1315 | |
| 1316 | def _add_biencoder_args(parser): |
| 1317 | group = parser.add_argument_group(title="biencoder") |
| 1318 | |
| 1319 | # network size |
| 1320 | group.add_argument( |
| 1321 | "--ict-head-size", |
| 1322 | type=int, |
| 1323 | default=None, |
| 1324 | help="Size of block embeddings to be used in ICT and " |
| 1325 | "REALM (paper default: 128)", |
| 1326 | ) |
| 1327 | group.add_argument( |
| 1328 | "--biencoder-projection-dim", |
| 1329 | type=int, |
| 1330 | default=0, |
| 1331 | help="Size of projection head used in biencoder (paper" " default: 128)", |
| 1332 | ) |
| 1333 | group.add_argument( |
| 1334 | "--biencoder-shared-query-context-model", |
| 1335 | action="store_true", |
| 1336 | help="Whether to share the parameters of the query " |
| 1337 | "and context models or not", |
| 1338 | ) |
| 1339 | |
| 1340 | # checkpointing |
| 1341 | group.add_argument( |
| 1342 | "--ict-load", |
| 1343 | type=str, |
| 1344 | default=None, |
| 1345 | help="Directory containing an ICTBertModel checkpoint", |
| 1346 | ) |
| 1347 | group.add_argument( |
| 1348 | "--bert-load", |
| 1349 | type=str, |
| 1350 | default=None, |
| 1351 | help="Directory containing an BertModel checkpoint " |
| 1352 | "(needed to start ICT and REALM)", |
| 1353 | ) |
| 1354 | |
| 1355 | # data |
| 1356 | group.add_argument( |
| 1357 | "--titles-data-path", |
| 1358 | type=str, |
| 1359 | default=None, |
| 1360 | help="Path to titles dataset used for ICT", |
| 1361 | ) |
| 1362 | group.add_argument( |
| 1363 | "--query-in-block-prob", |
| 1364 | type=float, |
| 1365 | default=0.1, |
| 1366 | help="Probability of keeping query in block for " "ICT dataset", |
| 1367 | ) |
| 1368 | group.add_argument( |
| 1369 | "--use-one-sent-docs", |
| 1370 | action="store_true", |
| 1371 | help="Whether to use one sentence documents in ICT", |
| 1372 | ) |
| 1373 | group.add_argument( |