Something went wrong. Try again.
This repository has no description
Something went wrong. Try again.
1.4 kB · 49 lines
Python
at main
1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950from typing import List, Optional
import pydanticimport numpy as np
# Global list mapping each dihedral transform id to its inverse.# Index corresponds to the original tid, and the value is its inverse.DIHEDRAL_INVERSE = [0, 3, 2, 1, 4, 5, 6, 7]
class PuzzleDatasetMetadata(pydantic.BaseModel): pad_id: int ignore_label_id: Optional[int] blank_identifier_id: int vocab_size: int seq_len: int num_puzzle_identifiers: int total_groups: int mean_puzzle_examples: float total_puzzles: int sets: List[str]
def dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray: """8 dihedral symmetries by rotate, flip and mirror""" if tid == 0: return arr # identity elif tid == 1: return np.rot90(arr, k=1) elif tid == 2: return np.rot90(arr, k=2) elif tid == 3: return np.rot90(arr, k=3) elif tid == 4: return np.fliplr(arr) # horizontal flip elif tid == 5: return np.flipud(arr) # vertical flip elif tid == 6: return arr.T # transpose (reflection along main diagonal) elif tid == 7: return np.fliplr(np.rot90(arr, k=1)) # anti-diagonal reflection else: return arr def inverse_dihedral_transform(arr: np.ndarray, tid: int) -> np.ndarray: return dihedral_transform(arr, DIHEDRAL_INVERSE[tid])