diff --git a/pyterrier_dr/prf.py b/pyterrier_dr/prf.py index 1004be3..6f72d47 100644 --- a/pyterrier_dr/prf.py +++ b/pyterrier_dr/prf.py @@ -79,6 +79,9 @@ def __init__(self, ): self.k = k + def compile(self) -> pt.Transformer: + return pt.RankCutoff(self.k) >> self + @pta.transform.by_query(add_ranks=False) def transform(self, inp: pd.DataFrame) -> pd.DataFrame: """Performs Average PRF on the input dataframe."""