diff options
Diffstat (limited to 'worldalign/tier0_pipeline.py')
| -rw-r--r-- | worldalign/tier0_pipeline.py | 37 |
1 files changed, 33 insertions, 4 deletions
diff --git a/worldalign/tier0_pipeline.py b/worldalign/tier0_pipeline.py index 778e20f..01c78d2 100644 --- a/worldalign/tier0_pipeline.py +++ b/worldalign/tier0_pipeline.py @@ -56,6 +56,12 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--seed", type=int, default=0) parser.add_argument("--output", required=True) parser.add_argument("--states-output", default="") + parser.add_argument( + "--text-omit", + default="", + help="Comma-separated factors the captions never state " + "(colour, count, size). Vision still sees them.", + ) return parser.parse_args() @@ -188,7 +194,17 @@ def factor_vector(colour: int, count: int, size: int, classes: int) -> torch.Ten return vector -def encode_caption(caption: str, colour_rank: dict[str, int], classes: int) -> torch.Tensor: +def encode_caption( + caption: str, colour_rank: dict[str, int], classes: int, omit: frozenset = frozenset() +) -> torch.Tensor: + """Encode a caption, optionally suppressing factors the text never states. + + Suppressing a factor makes the caption describe strictly less of the scene + while vision continues to see all of it, which is the controlled form of + the situation on photographs: a region description and a patch descriptor + overlap on some aspects and not others. It is the intervention that tests + whether corpus overlap is what sets the width of the shared spectrum. + """ vectors = [] for phrase in parse_group_phrases(caption): tokens = phrase.split() @@ -201,10 +217,20 @@ def encode_caption(caption: str, colour_rank: dict[str, int], classes: int) -> t ) ) size = 0 if "small" in tokens else (2 if "large" in tokens else 1) - vectors.append(factor_vector(colour, count, size, classes)) + vector = factor_vector(colour, count, size, classes) + if "colour" in omit: + vector[:classes] = 0.0 + if "count" in omit: + vector[classes : classes + 4] = 0.0 + if "size" in omit: + vector[classes + 4 :] = 0.0 + vectors.append(vector) if not vectors: vectors = [factor_vector(0, 1, 1, classes)] - return F.normalize(torch.stack(vectors), dim=-1) + stacked = torch.stack(vectors) + if stacked.abs().sum() == 0: + stacked = stacked + 1.0 + return F.normalize(stacked, dim=-1) def moment_state(states: torch.Tensor) -> torch.Tensor: @@ -264,7 +290,10 @@ def main() -> None: view_states = [moment_state(item) for item in sets] visual_field = torch.stack(per_view_fields).mean(0) - text_sets = [encode_caption(captions[row][0], colour_rank, classes) for row in rows] + omit = frozenset(f for f in args.text_omit.split(",") if f) + text_sets = [ + encode_caption(captions[row][0], colour_rank, classes, omit) for row in rows + ] text_field = moment_field(text_sets) mask = ~np.eye(len(rows), dtype=bool) |