function useCollatedLogitData()

in neuron_viewer/src/TransformerDebugger/cards/LogitsDisplay.tsx [219:256]


function useCollatedLogitData(
  rightResponseData: MultipleTopKDerivedScalarsResponseData | null,
  leftResponseData: MultipleTopKDerivedScalarsResponseData
): LogitDatum[] {
  return useMemo(() => {
    const lookupTable = joinIndices(
      rightResponseData?.nodeIndices ?? [],
      leftResponseData.nodeIndices
    );
    const collatedLogitData: LogitDatum[] = lookupTable.nodeIndices.map((nodeIndex, i) => {
      const rightIndex = lookupTable.rightArrayIndices[i];
      const leftIndex = lookupTable.leftArrayIndices[i];
      const rightLogit =
        rightIndex === undefined || !rightResponseData
          ? undefined
          : rightResponseData.activationsByGroupId[GroupId.LOGITS][rightIndex];
      const leftLogit =
        leftIndex === undefined
          ? undefined
          : leftResponseData.activationsByGroupId[GroupId.LOGITS][leftIndex];
      const diffLogit = diffOptionalNumbers(leftLogit, rightLogit);
      let token =
        leftIndex === undefined || !leftResponseData
          ? undefined
          : leftResponseData.vocabTokenStringsForIndices![leftIndex];
      if (token === undefined && rightIndex !== undefined && rightResponseData) {
        token = rightResponseData?.vocabTokenStringsForIndices![rightIndex];
      }
      return {
        token: token ?? "<missing>",
        rightLogit,
        leftLogit,
        diffLogit,
      };
    });
    return collatedLogitData;
  }, [rightResponseData, leftResponseData]);
}