Update app.py

app.py

@@ -15,7 +15,7 @@
 __author__ = 'Dmitry Ustalov'
 __license__ = 'Apache 2.0'
 
-from typing import IO, Tuple
+from typing import IO, Tuple, List, cast
 
 import gradio as gr
 import numpy as np
@@ -32,11 +32,44 @@ def visualize(df_pairwise: pd.DataFrame) -> Figure:
     return fig
 
 
-# https://gist.github.com/dustalov/41678b70c40ba5a55430fa5e77b121d9#file-newman-py
-def aggregate(wins: npt.NDArray[np.int64], ties: npt.NDArray[np.int64],
-              seed: int = 0, tolerance: float = 10e-6, limit: int = 20) -> npt.NDArray[np.float64]:
-    assert wins.shape == ties.shape, 'wins and ties shapes are different'
+# https://gist.github.com/dustalov/41678b70c40ba5a55430fa5e77b121d9#file-bradley_terry-py
+def bradley_terry(wins: npt.NDArray[np.int64], ties: npt.NDArray[np.int64],
+                  seed: int = 0, tolerance: float = 10e-6, limit: int = 20) -> npt.NDArray[np.float64]:
+    M = wins + .5 * ties
+
+    T = M.T + M
+    active = T > 0
+
+    w = M.sum(axis=1)
+
+    Z = np.zeros_like(M, dtype=float)
+
+    p = np.ones(M.shape[0])
+    p_new = p.copy()
+
+    converged, iterations = False, 0
+
+    while not converged:
+        iterations += 1
+
+        P = np.broadcast_to(p, M.shape)
+
+        Z[active] = T[active] / (P[active] + P.T[active])
 
+        p_new[:] = w
+        p_new /= Z.sum(axis=0)
+        p_new /= p_new.sum()
+
+        converged = bool(np.linalg.norm(p_new - p) < tolerance) or (iterations >= limit)
+
+        p[:] = p_new
+
+    return p
+
+
+# https://gist.github.com/dustalov/41678b70c40ba5a55430fa5e77b121d9#file-newman-py
+def newman(wins: npt.NDArray[np.int64], ties: npt.NDArray[np.int64],
+           seed: int = 0, tolerance: float = 10e-6, limit: int = 20) -> npt.NDArray[np.float64]:
     rng = np.random.default_rng(seed)
 
     pi, v = rng.random(wins.shape[0]), rng.random()
@@ -82,10 +115,19 @@ def aggregate(wins: npt.NDArray[np.int64], ties: npt.NDArray[np.int64],
     return pi
 
 
-def handler(file: IO[bytes], seed: int) -> Tuple[pd.DataFrame, Figure]:
+ALGORITHMS = {
+    'Bradley-Terry (1952)': bradley_terry,
+    'Newman (2023)': newman,
+}
+
+
+def handler(file: IO[bytes], algorithm: str, seed: int) -> Tuple[pd.DataFrame, Figure]:
     if file is None:
         raise gr.Error('File must be uploaded')
 
+    if algorithm not in ALGORITHMS:
+        raise gr.Error(f'Unknown algorithm: {algorithm}')
+
     try:
         df = pd.read_csv(file.name, dtype=str)
     except ValueError as e:
@@ -117,7 +159,9 @@ def handler(file: IO[bytes], seed: int) -> Tuple[pd.DataFrame, Figure]:
     ties = df_ties.to_numpy(dtype=np.int64)
     ties += ties.T
 
-    scores = aggregate(wins, ties, seed=seed)
+    assert wins.shape == ties.shape, 'wins and ties shapes are different'
+
+    scores = ALGORITHMS[algorithm](wins, ties, seed=seed)
 
     df_result = pd.DataFrame(data={'score': scores}, index=index)
     df_result['rank'] = df_result['score'].rank(na_option='bottom', ascending=False).astype(int)
@@ -142,6 +186,11 @@ def main() -> None:
                 file_types=['.tsv', '.csv'],
                 label='Comparisons'
             ),
+            gr.Dropdown(
+                choices=cast(List[str], ALGORITHMS),
+                value='Bradley-Terry (1952)',
+                label='Algorithm'
+            ),
             gr.Number(
                 label='Seed',
                 precision=0
@@ -157,7 +206,8 @@ def main() -> None:
             )
         ],
         examples=[
-            ['food.csv', 0]
+            ['food.csv', 'Bradley-Terry (1952)', 0],
+            ['food.csv', 'Newman (2023)', 0]
         ],
         title='Pair2Rank: Turn Your Side-by-Side Comparisons into Ranking!',
         description='''