Skip to content

Commit 424a2dc

Browse files
jackzhxngjackzhxng
authored
Add thinking mode toggle and UX improvements for Qwen3 on iOS app (pytorch#10614)
1 parent 3cf9863 commit 424a2dc

File tree

1 file changed

+149
-82
lines changed

1 file changed

+149
-82
lines changed

‎examples/demo-apps/apple_ios/LLaMA/LLaMA/Application/ContentView.swift‎

Lines changed: 149 additions & 82 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,8 @@ struct ContentView: View {
6262
@State private var isGenerating = false
6363
@State private var shouldStopGenerating = false
6464
@State private var shouldStopShowingToken = false
65+
@State private var thinkingMode = false
66+
@State private var showThinkingModeNotification = false
6567
private let runnerQueue = DispatchQueue(label: "org.pytorch.executorch.llama")
6668
@StateObject private var runnerHolder = RunnerHolder()
6769
@StateObject private var resourceManager = ResourceManager()
@@ -119,107 +121,136 @@ struct ContentView: View {
119121

120122
var body: some View {
121123
NavigationView {
122-
VStack {
123-
if showingSettings {
124-
VStack(spacing: 20) {
125-
HStack {
126-
VStack(spacing: 10) {
127-
Button(action: { pickerType = .model }) {
128-
Label(modelTitle, systemImage: "doc")
129-
.lineLimit(1)
130-
.truncationMode(.middle)
131-
.frame(maxWidth: 300, alignment: .leading)
132-
}
133-
Button(action: { pickerType = .tokenizer }) {
134-
Label(tokenizerTitle, systemImage: "doc")
135-
.lineLimit(1)
136-
.truncationMode(.middle)
137-
.frame(maxWidth: 300, alignment: .leading)
124+
ZStack {
125+
VStack {
126+
if showingSettings {
127+
VStack(spacing: 20) {
128+
HStack {
129+
VStack(spacing: 10) {
130+
Button(action: { pickerType = .model }) {
131+
Label(modelTitle, systemImage: "doc")
132+
.lineLimit(1)
133+
.truncationMode(.middle)
134+
.frame(maxWidth: 300, alignment: .leading)
135+
}
136+
Button(action: { pickerType = .tokenizer }) {
137+
Label(tokenizerTitle, systemImage: "doc")
138+
.lineLimit(1)
139+
.truncationMode(.middle)
140+
.frame(maxWidth: 300, alignment: .leading)
141+
}
138142
}
143+
.padding()
144+
.background(Color.gray.opacity(0.1))
145+
.cornerRadius(10)
146+
.fixedSize(horizontal: true, vertical: false)
147+
Spacer()
139148
}
140149
.padding()
141-
.background(Color.gray.opacity(0.1))
142-
.cornerRadius(10)
143-
.fixedSize(horizontal: true, vertical: false)
144-
Spacer()
145150
}
146-
.padding()
147151
}
148-
}
149152

150-
MessageListView(messages: $messages)
151-
.simultaneousGesture(
152-
DragGesture().onChanged { value in
153-
if value.translation.height > 10 {
154-
hideKeyboard()
153+
MessageListView(messages: $messages)
154+
.simultaneousGesture(
155+
DragGesture().onChanged { value in
156+
if value.translation.height > 10 {
157+
hideKeyboard()
158+
}
159+
showingSettings = false
160+
textFieldFocused = false
155161
}
162+
)
163+
.onTapGesture {
156164
showingSettings = false
157165
textFieldFocused = false
158166
}
159-
)
160-
.onTapGesture {
161-
showingSettings = false
162-
textFieldFocused = false
163-
}
164167

165-
HStack {
166-
Button(action: {
167-
imagePickerSourceType = .photoLibrary
168-
isImagePickerPresented = true
169-
}) {
170-
Image(systemName: "photo.on.rectangle")
171-
.resizable()
172-
.scaledToFit()
173-
.frame(width: 24, height: 24)
174-
}
175-
.background(Color.clear)
176-
.cornerRadius(8)
177-
178-
Button(action: {
179-
if UIImagePickerController.isSourceTypeAvailable(.camera) {
180-
imagePickerSourceType = .camera
168+
HStack {
169+
Button(action: {
170+
imagePickerSourceType = .photoLibrary
181171
isImagePickerPresented = true
182-
} else {
183-
print("Camera not available")
172+
}) {
173+
Image(systemName: "photo.on.rectangle")
174+
.resizable()
175+
.scaledToFit()
176+
.frame(width: 24, height: 24)
184177
}
185-
}) {
186-
Image(systemName: "camera")
187-
.resizable()
188-
.scaledToFit()
189-
.frame(width: 24, height: 24)
190-
}
191-
.background(Color.clear)
192-
.cornerRadius(8)
178+
.background(Color.clear)
179+
.cornerRadius(8)
193180

194-
TextField(placeholder, text: $prompt, axis: .vertical)
195-
.padding(8)
196-
.background(Color.gray.opacity(0.1))
197-
.cornerRadius(20)
198-
.lineLimit(1...10)
199-
.overlay(
200-
RoundedRectangle(cornerRadius: 20)
201-
.stroke(isInputEnabled ? Color.blue : Color.gray, lineWidth: 1)
202-
)
203-
.disabled(!isInputEnabled)
204-
.focused($textFieldFocused)
205-
.onAppear { textFieldFocused = false }
206-
.onTapGesture {
207-
showingSettings = false
181+
Button(action: {
182+
if UIImagePickerController.isSourceTypeAvailable(.camera) {
183+
imagePickerSourceType = .camera
184+
isImagePickerPresented = true
185+
} else {
186+
print("Camera not available")
187+
}
188+
}) {
189+
Image(systemName: "camera")
190+
.resizable()
191+
.scaledToFit()
192+
.frame(width: 24, height: 24)
193+
}
194+
.background(Color.clear)
195+
.cornerRadius(8)
196+
197+
if resourceManager.isModelValid && ModelType.fromPath(resourceManager.modelPath) == .qwen3 {
198+
Button(action: {
199+
thinkingMode.toggle()
200+
showThinkingModeNotification = true
201+
DispatchQueue.main.asyncAfter(deadline: .now() + 3) {
202+
showThinkingModeNotification = false
203+
}
204+
}) {
205+
Image(systemName: "brain")
206+
.resizable()
207+
.scaledToFit()
208+
.frame(width: 24, height: 24)
209+
.foregroundColor(thinkingMode ? .blue : .gray)
210+
}
211+
.background(Color.clear)
212+
.cornerRadius(8)
208213
}
209214

210-
Button(action: isGenerating ? stop : generate) {
211-
Image(systemName: isGenerating ? "stop.circle" : "arrowshape.up.circle.fill")
212-
.resizable()
213-
.aspectRatio(contentMode: .fit)
214-
.frame(height: 28)
215+
TextField(placeholder, text: $prompt, axis: .vertical)
216+
.padding(8)
217+
.background(Color.gray.opacity(0.1))
218+
.cornerRadius(20)
219+
.lineLimit(1...10)
220+
.overlay(
221+
RoundedRectangle(cornerRadius: 20)
222+
.stroke(isInputEnabled ? Color.blue : Color.gray, lineWidth: 1)
223+
)
224+
.disabled(!isInputEnabled)
225+
.focused($textFieldFocused)
226+
.onAppear { textFieldFocused = false }
227+
.onTapGesture {
228+
showingSettings = false
229+
}
230+
231+
Button(action: isGenerating ? stop : generate) {
232+
Image(systemName: isGenerating ? "stop.circle" : "arrowshape.up.circle.fill")
233+
.resizable()
234+
.aspectRatio(contentMode: .fit)
235+
.frame(height: 28)
236+
}
237+
.disabled(isGenerating ? shouldStopGenerating : (!isInputEnabled || prompt.isEmpty))
215238
}
216-
.disabled(isGenerating ? shouldStopGenerating : (!isInputEnabled || prompt.isEmpty))
239+
.padding([.leading, .trailing, .bottom], 10)
217240
}
218-
.padding([.leading, .trailing, .bottom], 10)
219241
.sheet(isPresented: $isImagePickerPresented, onDismiss: addSelectedImageMessage) {
220242
ImagePicker(selectedImage: $selectedImage, sourceType: imagePickerSourceType)
221243
.id(imagePickerSourceType.rawValue)
222244
}
245+
246+
if showThinkingModeNotification {
247+
Text(thinkingMode ? "Thinking mode enabled" : "Thinking mode disabled")
248+
.padding(8)
249+
.background(Color(UIColor.secondarySystemBackground))
250+
.cornerRadius(8)
251+
.transition(.opacity)
252+
.animation(.easeInOut(duration: 0.2), value: showThinkingModeNotification)
253+
}
223254
}
224255
.navigationBarTitle(title, displayMode: .inline)
225256
.navigationBarItems(
@@ -435,7 +466,10 @@ struct ContentView: View {
435466
let prompt: String
436467
switch modelType {
437468
case .qwen3:
438-
prompt = String(format: Constants.qwen3PromptTemplate, text)
469+
let basePrompt = String(format: Constants.qwen3PromptTemplate, text)
470+
// If thinking mode is enabled for Qwen, don't skip the <think></think> special tokens
471+
// and have them be generated.
472+
prompt = thinkingMode ? basePrompt.replacingOccurrences(of: "<think>\n\n</think>\n\n\n", with: "") : basePrompt
439473
case .llama:
440474
prompt = String(format: Constants.llama3PromptTemplate, text)
441475
case .llava:
@@ -445,12 +479,45 @@ struct ContentView: View {
445479
try runnerHolder.runner?.generate(prompt, sequenceLength: seq_len) { token in
446480

447481
if token != prompt {
448-
// hack to fix the issue that extension/llm/runner/text_token_generator.h
449-
// keeps generating after <|eot_id|>
450482
if token == "<|eot_id|>" {
483+
// hack to fix the issue that extension/llm/runner/text_token_generator.h
484+
// keeps generating after <|eot_id|>
451485
shouldStopShowingToken = true
486+
} else if token == "<|im_end|>" {
487+
// Qwen3 specific token.
488+
// Skip.
489+
} else if token == "<think>" {
490+
// Qwen3 specific token.
491+
let textToFlush = tokens.joined()
492+
let flushedTokenCount = tokens.count
493+
tokens = []
494+
DispatchQueue.main.async {
495+
var message = messages.removeLast()
496+
message.text += textToFlush
497+
message.text += message.text.isEmpty ? "Thinking...\n\n" : "\n\nThinking...\n\n"
498+
message.format = .italic
499+
message.tokenCount += flushedTokenCount + 1 // + 1 for the start thinking token.
500+
message.dateUpdated = Date()
501+
messages.append(message)
502+
}
503+
} else if token == "</think>" {
504+
// Qwen3 specific token.
505+
let textToFlush = tokens.joined()
506+
let flushedTokenCount = tokens.count
507+
tokens = []
508+
DispatchQueue.main.async {
509+
var message = messages.removeLast()
510+
message.text += textToFlush
511+
message.text += "\n\nFinished thinking.\n\n"
512+
message.format = .italic
513+
message.tokenCount += flushedTokenCount + 1 // + 1 for the end thinking token.
514+
message.dateUpdated = Date()
515+
messages.append(message)
516+
}
452517
} else {
453518
tokens.append(token.trimmingCharacters(in: .newlines))
519+
// Flush tokens in groups of 3 so that it's closer to whole words being generated
520+
// rather than parts of words (tokens).
454521
if tokens.count > 2 {
455522
let text = tokens.joined()
456523
let count = tokens.count

0 commit comments

Comments
 (0)