Skip to content

Commit f1d2c71

Browse files
hietalajuliushietalajulius
committed
fix: add model load time and tokens per second info
1 parent 9ef8ed2 commit f1d2c71

File tree

4 files changed

+160
-108
lines changed

4 files changed

+160
-108
lines changed

‎examples/demo-apps/react-native/rnllama/README.md‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -40,5 +40,4 @@ A React Native mobile application for running LLaMA language models using ExecuT
4040

4141
10. Select the model and tokenizer in the app to start chatting:
4242

43-
[![rnllama]](https://github.com/user-attachments/assets/0ca5a258-a248-4247-9e09-451945c1014b)
44-
43+
[![rnllama]](https://github.com/user-attachments/assets/0ca5a258-a248-4247-9e09-451945c1014b)

‎examples/demo-apps/react-native/rnllama/app/index.tsx‎

Lines changed: 157 additions & 102 deletions
Original file line numberDiff line numberDiff line change
@@ -21,48 +21,76 @@ import { Ionicons } from '@expo/vector-icons';
2121
export default function Index() {
2222
const [prompt, setPrompt] = useState('');
2323
const [currentOutput, setCurrentOutput] = useState('');
24+
2425
const [isGenerating, setIsGenerating] = useState(false);
25-
const [isStopped, setIsStopped] = useState(false);
26+
2627
const [modelPath, setModelPath] = useState('');
2728
const [modelName, setModelName] = useState('');
2829
const [tokenizerPath, setTokenizerPath] = useState('');
2930
const [tokenizerName, setTokenizerName] = useState('');
31+
3032
const [isInitialized, setIsInitialized] = useState(false);
3133
const [isInitializing, setIsInitializing] = useState(false);
32-
const [history, setHistory] = useState<Array<{ input: boolean, text: string }>>([]);
34+
35+
const [history, setHistory] = useState<Array<{
36+
input: boolean, text: string, stats?: {
37+
tokens: number;
38+
time: number;
39+
};
40+
}>>([]);
41+
42+
const [modelLoadTime, setModelLoadTime] = useState<number | null>(null);
43+
44+
const [currentGenerationStartTime, setCurrentGenerationStartTime] = useState<number | null>(null);
45+
const [currentNumTokens, setCurrentNumTokens] = useState(0);
46+
3347
const scrollViewRef = useRef();
3448

49+
const handleGenerationStopped = () => {
50+
LLaMABridge.stop();
51+
const generationEndTime = Date.now();
52+
const stats = currentGenerationStartTime !== null ? {
53+
tokens: currentNumTokens,
54+
time: generationEndTime - currentGenerationStartTime
55+
} : undefined;
56+
57+
setHistory(prevHistory => [...prevHistory, { input: false, text: currentOutput.trim(), stats }]);
58+
setIsGenerating(false);
59+
setCurrentOutput('');
60+
setCurrentNumTokens(0);
61+
}
62+
3563
useEffect(() => {
3664
const unsubscribe = LLaMABridge.onToken((token) => {
37-
if (!isStopped) {
65+
if (isGenerating) {
3866
// Natural stop
3967
if (token === "<|eot_id|>") {
40-
setIsGenerating(false);
41-
setCurrentOutput(prev => {
42-
if (prev.trim()) {
43-
setHistory(prevHistory => [...prevHistory, { input: false, text: prev.trim() }]);
44-
}
45-
return '';
46-
});
68+
handleGenerationStopped();
4769
return;
4870
}
49-
71+
5072
// Skip template tokens
5173
if (token === formatPrompt('') ||
52-
token.includes("<|begin_of_text|>") ||
53-
token.includes("<|start_header_id|>") ||
54-
token.includes("<|end_header_id|>") ||
55-
token.includes("assistant")) {
74+
token.includes("<|begin_of_text|>") ||
75+
token.includes("<|start_header_id|>") ||
76+
token.includes("<|end_header_id|>") ||
77+
token.includes("assistant")) {
5678
return;
5779
}
58-
80+
5981
// Add token without leading newlines
60-
setCurrentOutput(prev => prev + token.replace(/^\n+/, ''));
82+
if (currentNumTokens === 0) {
83+
setCurrentOutput(prev => prev + token.replace(/^\n+/, ''));
84+
} else {
85+
setCurrentOutput(prev => prev + token);
86+
}
87+
88+
setCurrentNumTokens(prev => prev + 1);
6189
}
6290
});
63-
91+
6492
return () => unsubscribe();
65-
}, [isStopped, currentOutput]);
93+
}, [isGenerating, currentNumTokens]);
6694

6795

6896
const formatPrompt = (text: string) => {
@@ -74,11 +102,11 @@ export default function Index() {
74102
return;
75103
}
76104

77-
setIsStopped(false);
105+
setCurrentGenerationStartTime(Date.now());
106+
78107
const newPrompt = prompt.trim();
79108
setPrompt('');
80109
setIsGenerating(true);
81-
setCurrentOutput('');
82110

83111
// Add the user message immediately
84112
const userMessage = { input: true, text: newPrompt };
@@ -94,18 +122,6 @@ export default function Index() {
94122
}
95123
};
96124

97-
const handleStop = () => {
98-
if (!isGenerating) return;
99-
100-
setIsStopped(true);
101-
LLaMABridge.stop();
102-
103-
if (currentOutput) {
104-
setHistory(prev => [...prev, { input: false, text: currentOutput.trim() }]);
105-
}
106-
setCurrentOutput('');
107-
setIsGenerating(false);
108-
};
109125

110126
const selectModel = async () => {
111127
try {
@@ -149,17 +165,21 @@ export default function Index() {
149165
setCurrentOutput('');
150166
return;
151167
}
152-
168+
153169
if (!modelPath || !tokenizerPath) {
154170
Alert.alert('Error', 'Please select both model and tokenizer files first');
155171
return;
156172
}
157-
173+
158174
setIsInitializing(true);
159175
try {
176+
const startTime = Date.now();
160177
await LLaMABridge.initialize(modelPath, tokenizerPath);
178+
const modelLoadTime = Date.now() - startTime;
179+
setModelLoadTime(modelLoadTime);
180+
161181
setIsInitialized(true);
162-
Alert.alert('Success', 'LLaMA initialized successfully');
182+
Alert.alert('Success', `Model loaded in ${(modelLoadTime / 1000).toFixed(1)}s`);
163183
} catch (error) {
164184
console.error('Failed to initialize LLaMA:', error);
165185
Alert.alert('Error', 'Failed to initialize LLaMA');
@@ -184,7 +204,7 @@ export default function Index() {
184204
<View style={styles.header}>
185205
<Text style={styles.headerTitle}>rnllama</Text>
186206
</View>
187-
207+
188208
<View style={styles.setupBar}>
189209
<View style={styles.setupControls}>
190210
<TouchableOpacity
@@ -206,75 +226,95 @@ export default function Index() {
206226
</Text>
207227
</TouchableOpacity>
208228
</View>
209-
<TouchableOpacity
210-
style={[
211-
styles.initButton,
212-
isInitialized ? styles.setupComplete : styles.setupIncomplete,
213-
(!modelPath || !tokenizerPath || isInitializing) && styles.buttonDisabled
214-
]}
215-
onPress={initializeLLaMA}
216-
disabled={!modelPath || !tokenizerPath || isInitializing}
217-
>
218-
{isInitializing ? (
219-
<ActivityIndicator size="small" color="#fff" />
220-
) : (
221-
<Ionicons
222-
name={isInitialized ? "checkmark-circle-outline" : "power-outline"}
223-
size={24}
224-
color="#fff"
225-
/>
226-
)}
227-
</TouchableOpacity>
228-
</View>
229+
<View style={styles.initContainer}>
229230

231+
<TouchableOpacity
232+
style={[
233+
styles.initButton,
234+
isInitialized ? styles.setupComplete : styles.setupIncomplete,
235+
(!modelPath || !tokenizerPath || isInitializing) && styles.buttonDisabled
236+
]}
237+
onPress={initializeLLaMA}
238+
disabled={!modelPath || !tokenizerPath || isInitializing}
239+
>
240+
{isInitializing ? (
241+
<ActivityIndicator size="small" color="#fff" />
242+
) : (
243+
<Ionicons
244+
name={isInitialized ? "checkmark-circle-outline" : "power-outline"}
245+
size={24}
246+
color="#fff"
247+
/>
248+
)}
249+
</TouchableOpacity>
250+
</View>
251+
</View>
252+
230253
<KeyboardAvoidingView
231254
behavior={Platform.OS === "ios" ? "padding" : "height"}
232255
style={styles.content}
233256
keyboardVerticalOffset={Platform.OS === "ios" ? 90 : 0}
234257
>
235258
<ScrollView
236-
ref={scrollViewRef}
237-
style={styles.chatContainer}
238-
contentContainerStyle={styles.chatContent}
239-
onContentSizeChange={() => scrollViewRef.current?.scrollToEnd({ animated: true })}
240-
>
241-
{!isInitialized ? (
242-
<View style={styles.initPrompt}>
243-
<Text style={styles.initPromptText}>
244-
Please select model and tokenizer files, then initialize LLaMA to begin chatting
245-
</Text>
246-
</View>
247-
) : history.length === 0 ? (
248-
<Pressable
249-
style={styles.emptyState}
250-
onLongPress={handleClearHistory}
251-
>
252-
<Text style={styles.emptyStateText}>Start a conversation</Text>
253-
<Text style={styles.emptyStateHint}>Long press to clear history</Text>
254-
</Pressable>
255-
) : (
256-
<Pressable onLongPress={handleClearHistory}>
257-
{history.map((message, index) => (
258-
<View
259-
key={index}
260-
style={[
261-
message.input ? styles.sentMessage : styles.receivedMessage
262-
]}
263-
>
264-
<Text style={message.input ? styles.sentMessageText : styles.receivedMessageText}>
265-
{message.text}
266-
</Text>
267-
</View>
268-
))}
269-
{currentOutput && (
270-
<View style={styles.receivedMessage}>
271-
<Text style={styles.receivedMessageText}>{currentOutput}</Text>
272-
</View>
273-
)}
274-
</Pressable>
275-
)}
276-
</ScrollView>
277-
259+
ref={scrollViewRef}
260+
style={styles.chatContainer}
261+
contentContainerStyle={styles.chatContent}
262+
onContentSizeChange={() => scrollViewRef.current?.scrollToEnd({ animated: true })}
263+
>
264+
{!isInitialized ? (
265+
<View style={styles.initPrompt}>
266+
<Text style={styles.initPromptText}>
267+
Please select model and tokenizer files, then initialize LLaMA to begin chatting
268+
</Text>
269+
</View>
270+
) : history.length === 0 ? (
271+
<Pressable
272+
style={styles.emptyState}
273+
onLongPress={handleClearHistory}
274+
>
275+
{modelLoadTime && (
276+
<Text style={styles.emptyStateText}>
277+
Model loading took {(modelLoadTime / 1000).toFixed(1)}s
278+
</Text>
279+
)}
280+
<Text style={styles.emptyStateText}>Start a conversation</Text>
281+
<Text style={styles.emptyStateHint}>Long press to clear history</Text>
282+
</Pressable>
283+
) : (
284+
<Pressable onLongPress={handleClearHistory}>
285+
{modelLoadTime && (
286+
<View style={styles.loadTimeContainer}>
287+
<Text style={styles.loadTimeMessage}>
288+
Model loading took {(modelLoadTime / 1000).toFixed(1)}s
289+
</Text>
290+
</View>
291+
)}
292+
{history.map((message, index) => (
293+
<View
294+
key={index}
295+
style={[
296+
message.input ? styles.sentMessage : styles.receivedMessage
297+
]}
298+
>
299+
<Text style={message.input ? styles.sentMessageText : styles.receivedMessageText}>
300+
{message.text}
301+
</Text>
302+
{!message.input && message.stats && (
303+
<Text style={styles.tokensPerSecondText}>
304+
{`Tokens/sec: ${(message.stats.tokens / (message.stats.time / 1000)).toFixed(2)}`}
305+
</Text>
306+
)}
307+
</View>
308+
))}
309+
{currentOutput && (
310+
<View style={styles.receivedMessage}>
311+
<Text style={styles.receivedMessageText}>{currentOutput}</Text>
312+
</View>
313+
)}
314+
</Pressable>
315+
)}
316+
</ScrollView>
317+
278318
<View style={styles.inputContainer}>
279319
<TextInput
280320
value={prompt}
@@ -287,8 +327,8 @@ export default function Index() {
287327
/>
288328
<TouchableOpacity
289329
style={[styles.sendButton, (!isInitialized || (!isGenerating && !prompt.trim())) && styles.buttonDisabled]}
290-
onPress={isGenerating ? handleStop : handleGenerate}
291-
disabled={!isInitialized || (!prompt.trim() && !isGenerating)} // This was backwards
330+
onPress={isGenerating ? handleGenerationStopped : handleGenerate}
331+
disabled={!isInitialized || (!prompt.trim() && !isGenerating)}
292332
>
293333
<Ionicons
294334
name={isGenerating ? "stop-outline" : "send-outline"}
@@ -307,11 +347,26 @@ const styles = StyleSheet.create({
307347
flex: 1,
308348
backgroundColor: '#000000',
309349
},
350+
loadTimeContainer: {
351+
alignItems: 'center',
352+
marginBottom: 16,
353+
padding: 8,
354+
backgroundColor: '#1A1A1A',
355+
borderRadius: 8,
356+
},
357+
loadTimeMessage: {
358+
color: '#666',
359+
fontSize: 14,
360+
},
361+
tokensPerSecondText: {
362+
color: '#666',
363+
fontSize: 12,
364+
marginTop: 4,
365+
},
310366
header: {
311367
padding: 16,
312368
borderBottomWidth: 1,
313369
borderBottomColor: '#333',
314-
alignItems: "start",
315370
},
316371
headerTitle: {
317372
fontSize: 24,

‎examples/demo-apps/react-native/rnllama/ios/LlamaBridge.mm‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -56,5 +56,4 @@ + (BOOL)requiresMainQueueSetup {
5656
});
5757
}
5858

59-
@end
60-
59+
@end

‎examples/demo-apps/react-native/rnllama/ios/rnllama/main.m‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -6,5 +6,4 @@ int main(int argc, char * argv[]) {
66
@autoreleasepool {
77
return UIApplicationMain(argc, argv, nil, NSStringFromClass([AppDelegate class]));
88
}
9-
}
10-
9+
}

0 commit comments

Comments
 (0)