Integrate maia3 model for analysis with 600-2600 rating range

.gitignore

@@ -37,6 +37,10 @@ yarn-error.log*
 # typescript
 *.tsbuildinfo
 
-# Large model weights - hosted on GitHub Releases, not committed to the repo
+# Large model weights
 /public/maia2/*.onnx
 /public/stockfish/*.nnue
+
+# Maia3 training code and checkpoints
+/maia3_code/
+*.pt

src/api/analysis.ts

@@ -226,7 +226,7 @@ export const fetchAnalyzedMaiaGame = async (
   const blackPlayer = data['black_player'] as Player
   const whitePlayer = data['white_player'] as Player
 
-  const maiaPattern = /maia_kdd_1\d00/
+  const maiaPattern = /maia_kdd_\d{3,4}/
 
   if (blackPlayer.name && maiaPattern.test(blackPlayer.name)) {
     blackPlayer.name = blackPlayer.name.replace('maia_kdd_', 'Maia ')

src/components/Analysis/BoardChrome.tsx

@@ -250,9 +250,9 @@ export const AnalysisMaiaWinrateBar: React.FC<AnalysisMaiaWinrateBarProps> = ({
     : 'w-4'
   const tickTextClass = isDesktop
     ? isExpandedDesktop
-      ? 'text-[9px]'
-      : 'text-[8px]'
-    : 'text-[8px]'
+      ? 'text-[8px]'
+      : 'text-[7px]'
+    : 'text-[7px]'
   const bubbleClass = isDesktop
     ? isExpandedDesktop
       ? 'h-6 min-w-[42px] px-2 text-[11px]'

src/components/Analysis/Highlight.tsx

@@ -355,7 +355,8 @@ export const Highlight: React.FC<Props> = ({
   }, [boardDescription?.segments?.length])
 
   const useCompactMobileColumnTitles = isMobile || simplified
-  const mobileMaiaColumnTitle = `Maia ${currentMaiaModel.slice(-4)}: Human Moves`
+  const maiaRating = currentMaiaModel.replace('maia_kdd_', '')
+  const mobileMaiaColumnTitle = `Maia ${maiaRating}: Human Moves`
   const mobileStockfishColumnTitle = 'SF 17: Engine Moves'
   const compactTitleRowClass = 'grid h-11 place-items-center'
   const splitTitleRowClass = 'grid h-12 place-items-center'
@@ -414,7 +415,7 @@ export const Highlight: React.FC<Props> = ({
               >
                 {useCompactMobileColumnTitles
                   ? mobileMaiaColumnTitle
-                  : `Maia ${currentMaiaModel.slice(-4)}`}
+                  : `Maia ${maiaRating}`}
               </div>
             ) : (
               <>
@@ -434,7 +435,7 @@ export const Highlight: React.FC<Props> = ({
                           key={model}
                           className="bg-transparent text-human-1"
                         >
-                          {`Maia ${model.slice(-4)}`}
+                          {`Maia ${model.replace('maia_kdd_', '')}`}
                         </option>
                       ))}
                     </select>
@@ -455,7 +456,7 @@ export const Highlight: React.FC<Props> = ({
                         key={model}
                         className="bg-transparent text-human-1"
                       >
-                        {`Maia ${model.slice(-4)}`}
+                        {`Maia ${model.replace('maia_kdd_', '')}`}
                       </option>
                     ))}
                   </select>

src/components/Analysis/MovesByRating.tsx

@@ -67,8 +67,11 @@ export const MovesByRating: React.FC<Props> = ({
             tickMargin={4}
             ticks={
               isMobile
-                ? [1100, 1300, 1500, 1700, 1900]
-                : [1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900]
+                ? [600, 1000, 1400, 1800, 2200, 2600]
+                : [
+                    600, 800, 1000, 1200, 1400, 1600, 1800, 2000, 2200, 2400,
+                    2600,
+                  ]
             }
           />
           <YAxis

src/components/Common/DownloadModelModal.tsx

@@ -103,7 +103,7 @@ export const DownloadModelModal: React.FC<Props> = ({
 
           <div className="flex flex-col gap-3 text-sm md:text-base">
             <p className="text-primary/80">
-              Our newest model, Maia 2, provides richer and more in-depth
+              Our newest model, Maia 3, provides richer and more in-depth
               analysis, allowing for:
             </p>
             <ul className="flex list-inside flex-col gap-1.5 pl-2 text-white/80">
@@ -120,7 +120,7 @@ export const DownloadModelModal: React.FC<Props> = ({
               </li>
             </ul>
             <p className="text-white/80">
-              Maia 2 runs entirely on your device and requires a one-time 90mb
+              Maia 3 runs entirely on your device and requires a one-time 87mb
               download.
             </p>
           </div>
@@ -174,8 +174,8 @@ export const DownloadModelModal: React.FC<Props> = ({
                   </span>
                 )}
                 <p className="font-medium">
-                  {isDownloading ? 'Downloading...' : 'Download Maia 2'}{' '}
-                  <span className="text-xs text-white/70">(90mb)</span>
+                  {isDownloading ? 'Downloading...' : 'Download Maia 3'}{' '}
+                  <span className="text-xs text-white/70">(87mb)</span>
                 </p>
               </div>
             </div>

src/components/Home/AboutMaia.tsx

@@ -237,9 +237,8 @@ export const AboutMaia = () => {
               Maia is trained to predict human moves rather than to find the
               optimal move in a position. As a result, Maia exhibits common
               human biases and makes many of the same mistakes that humans make.
-              We have trained a set of nine neural network engines, each
-              targeting a specific rating level on the Lichess.org rating scale,
-              from 1100 to 1900.
+              We have trained a neural network engine that can target any rating
+              level on the Lichess.org rating scale, from 600 to 2600.
             </p>
 
             <p className="text-primary/90">

src/components/Openings/MaiaRatingInsights.tsx

@@ -67,7 +67,7 @@ export const MaiaRatingInsights: React.FC<Props> = ({ ratingPrediction }) => {
     const upperBound = Math.round(
       ratingPrediction.predictedRating + ratingPrediction.standardDeviation,
     )
-    return `${Math.max(1100, lowerBound)} → ${Math.min(1900, upperBound)}`
+    return `${Math.max(600, lowerBound)} → ${Math.min(2600, upperBound)}`
   }
 
   return (
@@ -160,8 +160,8 @@ export const MaiaRatingInsights: React.FC<Props> = ({ ratingPrediction }) => {
           <span className="font-semibold text-human-2">
             How it&apos;s calculated:{' '}
           </span>
-          We analyze each position before your moves using 9 different Maia
-          models (1100-1900). For each model, we calculate the probability it
+          We analyze each position before your moves using 21 different Maia
+          models (600-2600). For each model, we calculate the probability it
           would play your chosen move. The final rating is the specific level
           that had the highest log likelihood of predicting your actual moves.
         </p>

src/constants/analysis.ts

@@ -18,6 +18,8 @@ export const MOVE_CLASSIFICATION_THRESHOLDS = {
 } as const
 
 export const DEFAULT_MAIA_MODEL = 'maia_kdd_1500' as const
+export const MAIA_MIN_RATING = 600 as const
+export const MAIA_MAX_RATING = 2600 as const
 export const MIN_STOCKFISH_DEPTH = 12 as const
 export const LEARN_FROM_MISTAKES_DEPTH = 12 as const
 export const STOCKFISH_DEBUG_RERUN_EVENT = 'maia:stockfish-debug-rerun' as const

src/constants/common.ts

@@ -1,14 +1,11 @@
-export const MAIA_MODELS = [
-  'maia_kdd_1100',
-  'maia_kdd_1200',
-  'maia_kdd_1300',
-  'maia_kdd_1400',
-  'maia_kdd_1500',
-  'maia_kdd_1600',
-  'maia_kdd_1700',
-  'maia_kdd_1800',
-  'maia_kdd_1900',
-]
+export const MAIA_MODELS = Array.from(
+  { length: 21 },
+  (_, i) => `maia_kdd_${600 + i * 100}`,
+)
+
+export const MAIA_RATINGS = MAIA_MODELS.map((m) =>
+  parseInt(m.replace('maia_kdd_', '')),
+)
 
 export const MAIA_MODELS_WITH_NAMES = MAIA_MODELS.map((model) => ({
   id: model,

src/contexts/MaiaEngineContext.tsx

@@ -36,8 +36,8 @@ export const MaiaEngineContextProvider: React.FC<{ children: ReactNode }> = ({
     const model = new Maia({
       model:
         process.env.NEXT_PUBLIC_MAIA_MODEL_URL ??
-        'https://raw.githubusercontent.com/CSSLab/maia-platform-frontend/e23a50e/public/maia2/maia_rapid.onnx',
-      modelVersion: process.env.NEXT_PUBLIC_MAIA_MODEL_VERSION ?? '1',
+        '/maia3/maia3_simplified.onnx',
+      modelVersion: process.env.NEXT_PUBLIC_MAIA_MODEL_VERSION ?? '2',
       setStatus: setStatus,
       setProgress: setProgress,
       setError: setError,

src/hooks/useAnalysisController/useEngineAnalysis.ts

@@ -1,7 +1,7 @@
 import { Chess } from 'chess.ts'
 import { fetchOpeningBookMoves } from 'src/api'
 import { useEffect, useContext, useRef, useState } from 'react'
-import { MAIA_MODELS } from 'src/constants/common'
+import { MAIA_MODELS, MAIA_RATINGS } from 'src/constants/common'
 import {
   STOCKFISH_DEBUG_RERUN_EVENT,
   STOCKFISH_DEBUG_RERUN_KEY,
@@ -58,10 +58,10 @@ export const useEngineAnalysis = (
       throw new Error('Maia engine not initialized')
     }
 
-    const { result } = await maia.maia.batchEvaluate(
-      Array(9).fill(board.fen()),
-      [1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900],
-      [1100, 1200, 1300, 1400, 1500, 1600, 1700, 1800, 1900],
+    const { result } = await maia.maia.batchEvaluateMaia3(
+      Array(MAIA_RATINGS.length).fill(board.fen()),
+      MAIA_RATINGS,
+      MAIA_RATINGS,
     )
 
     const maiaEvaluations: { [key: string]: MaiaEvaluation } = {}

src/hooks/useAnalysisController/useMoveRecommendations.ts

@@ -99,7 +99,7 @@ export const useMoveRecommendations = (
     const data = []
     for (const rating of MAIA_MODELS) {
       const entry: { [key: string]: number } = {
-        rating: parseInt(rating.slice(-4)),
+        rating: parseInt(rating.replace('maia_kdd_', '')),
       }
 
       for (const move of candidates) {