Skip to content

Commit

Permalink
[FR-22] feature: provide token-related information (#2995)
Browse files Browse the repository at this point in the history
# Add token metrics to LLM Chat UI

> [!NOTE]
> Many LLM models send token info as a NaN; So calculate usages on a front part.

Adds real-time performance metrics to the chat interface including:
- Token generation speed (tokens/second)
- Total tokens generated

The metrics are displayed in a tag at the bottom of the chat window, providing users visibility into the chat performance.

![image.png](https://graphite-user-uploaded-assets-prod.s3.amazonaws.com/2HueYSdFvL8pOB5mgrUQ/bbfdb740-c30f-4889-9414-520dab1847e8.png)

**Checklist:**
- [ ] Documentation
- [ ] Test case: Verify metrics update correctly when:
  - Messages are generated
  - Files are uploaded
  - Chat session continues over time
  • Loading branch information
agatha197 committed Jan 6, 2025
1 parent 34f9b24 commit e02fd86
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 2 deletions.
1 change: 1 addition & 0 deletions react/src/components/lablupTalkativotUI/ChatMessage.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ const ChatMessage: React.FC<{
message.role !== 'user'
? token.colorBgContainer
: token.colorBgContainerDisabled,
maxWidth: '100%',
}}
>
<ChatMessageContent>
Expand Down
45 changes: 43 additions & 2 deletions react/src/components/lablupTalkativotUI/LLMChatCard.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import {
DeleteOutlined,
LinkOutlined,
MoreOutlined,
RocketOutlined,
} from '@ant-design/icons';
import { Attachments, AttachmentsProps, Sender } from '@ant-design/x';
import { useControllableValue } from 'ahooks';
Expand All @@ -28,11 +29,13 @@ import {
FormInstance,
Input,
MenuProps,
Tag,
theme,
Typography,
} from 'antd';
import _ from 'lodash';
import { Scale } from 'lucide-react';
import React, { useEffect, useRef, useState } from 'react';
import React, { useEffect, useRef, useState, useMemo } from 'react';
import { useTranslation } from 'react-i18next';

export type BAIModel = {
Expand Down Expand Up @@ -137,6 +140,7 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
),
messages: body?.messages,
});
setStartTime(Date.now());
return result.toDataStreamResponse();
} else {
return fetch(input, init);
Expand All @@ -146,6 +150,8 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
const { token } = theme.useToken();
const { t } = useTranslation();

const [startTime, setStartTime] = useState<number | null>(null);

// If the `inputMessage` prop exists, the `input` state has to follow it.
useEffect(() => {
if (!_.isUndefined(inputMessage)) {
Expand All @@ -163,6 +169,24 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [submitKey]);

const totalTokens = useMemo(() => {
return (
_.sumBy(messages, (message) => message?.content.length) + input.length
);
}, [messages, input]);

const tokenPerSecond = useMemo(() => {
if (!_.isEmpty(messages) && startTime) {
const lastMessage = _.last(messages);
if (lastMessage?.role === 'assistant') {
const lastToken = lastMessage.content?.length || 0;
const elapsedTime = (Date.now() - startTime) / 1000;
return lastToken / elapsedTime;
}
}
return 0;
}, [messages, startTime]);

const [files, setFiles] = useState<AttachmentsProps['items']>([]);

const items: MenuProps['items'] = filterEmptyItem([
Expand Down Expand Up @@ -340,7 +364,9 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
// Filter after converting to `File`
const fileListArray = _.filter(fileList, Boolean);
const dataTransfer = new DataTransfer();
_.forEach(fileListArray, (file) => dataTransfer.items.add(file));
_.forEach(fileListArray, (file) => {
dataTransfer.items.add(file);
});

append(
{
Expand Down Expand Up @@ -433,6 +459,21 @@ const LLMChatCard: React.FC<LLMChatCardProps> = ({
/>
) : null}
<VirtualChatMessageList messages={messages} isStreaming={isLoading} />
<Flex justify="end" align="end" style={{ margin: token.marginSM }}>
<Tag>
<Flex gap={'xs'}>
<RocketOutlined />
<Flex gap={'xxs'}>
<Typography.Text>{tokenPerSecond.toFixed(2)}</Typography.Text>
<Typography.Text type="secondary">tok/s</Typography.Text>
</Flex>
<Flex gap={'xxs'}>
<Typography.Text>{totalTokens}</Typography.Text>
<Typography.Text type="secondary">total tokens</Typography.Text>
</Flex>
</Flex>
</Tag>
</Flex>
</Card>
);
};
Expand Down

0 comments on commit e02fd86

Please sign in to comment.