Skip to content

Commit

Permalink
FIX: Failed to translate if input text is long.
Browse files Browse the repository at this point in the history
- Increase number of tokens.
- Truncate tokens from the input and always gives suffix tokens.
  • Loading branch information
niw committed Sep 22, 2024
1 parent 19f92e5 commit 3748999
Show file tree
Hide file tree
Showing 4 changed files with 32 additions and 16 deletions.
20 changes: 14 additions & 6 deletions Sources/LlamaModel/LlamaContext.swift
Original file line number Diff line number Diff line change
Expand Up @@ -58,7 +58,7 @@ actor LlamaContext {
/// This variable is used to store temporarily invalid cchars
private var temporary_invalid_cchars: [CChar]

var n_len: Int32 = 1024
var n_len: Int32 = 4096
var n_cur: Int32 = 0

var n_decode: Int32 = 0
Expand All @@ -67,7 +67,7 @@ actor LlamaContext {
self.model = model
self.context = context
self.tokens_list = []
self.batch = llama_batch_init(512, 0, 1)
self.batch = llama_batch_init(2048, 0, 1)
self.temporary_invalid_cchars = []
let sparams = llama_sampler_chain_default_params()
self.sampling = llama_sampler_chain_init(sparams)
Expand Down Expand Up @@ -131,14 +131,22 @@ actor LlamaContext {
return batch.n_tokens;
}

func completion_init(text: String) throws {
tokens_list = tokenize(text: text, add_bos: true)
func completion_init(text: String, suffix: String? = nil) throws {
// See `llama_batch_init()` call about the max tokens in batch.
let text_tokens_list = tokenize(text: text, add_bos: true)
if let suffix {
let suffix_tokens_list = tokenize(text: suffix, add_bos: false)
let limit = max(0, 2048 - suffix_tokens_list.count)
tokens_list = Array(text_tokens_list.prefix(limit)) + suffix_tokens_list
} else {
tokens_list = Array(text_tokens_list.prefix(2048))
}

temporary_invalid_cchars = []

llama_batch_clear(&batch)

// See `llama_batch_init()`.
for i1 in 0..<min(512, tokens_list.count) {
for i1 in 0..<tokens_list.count {
let i = Int(i1)
llama_batch_add(&batch, tokens_list[i], Int32(i), [0], false)
}
Expand Down
4 changes: 2 additions & 2 deletions Sources/LlamaModel/LlamaModel.swift
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,13 @@ public struct LlamaModel: Sendable {
}.value
}

public func complete(_ prompt: String) -> AsyncThrowingStream<String, any Swift.Error> {
public func complete(_ prompt: String, suffix: String? = nil) -> AsyncThrowingStream<String, any Swift.Error> {
AsyncThrowingStream { continuation in
let task = Task {
do {
try Task.checkCancellation()

try await context.completion_init(text: prompt)
try await context.completion_init(text: prompt, suffix: suffix)

try Task.checkCancellation()

Expand Down
22 changes: 15 additions & 7 deletions Sources/TranslatorSupport/Translator.swift
Original file line number Diff line number Diff line change
Expand Up @@ -30,7 +30,7 @@ public enum Translator {
}

// See <https://huggingface.co/webbigdata/C3TR-Adapter_gguf>
private static let template =
private static let textFormat =
"""
You are a highly skilled professional Japanese-English and English-Japanese translator. Translate the given text accurately, taking into account the context and specific instructions provided. Steps may include hints enclosed in square brackets [] with the key and value separated by a colon:. Only when the subject is specified in the Japanese sentence, the subject will be added when translating into English. If no additional instructions or context are provided, use your expertise to consider what the most appropriate context is and provide a natural translation that aligns with that context. When translating, strive to faithfully reflect the meaning and tone of the original text, pay attention to cultural nuances and differences in language usage, and ensure that the translation is grammatically correct and easy to read. After completing the translation, review it once more to check for errors or unnatural expressions. For technical terms and proper nouns, either leave them in the original language or use appropriate translations as necessary. Take a deep breath, calm down, and start translating.
Expand All @@ -41,17 +41,25 @@ public enum Translator {
### Input:
%@
"""

private static let suffix =
"""
<end_of_turn>
<start_of_turn>### Response:
"""

static func prompt(mode: Mode, style: Style, input: String) -> String {
return String(
format: template,
mode.value(for: input),
style.rawValue,
input
struct Prompt {
var text: String
var suffix: String
}

static func prompt(mode: Mode, style: Style, input: String) -> Prompt {
Prompt(
text: String(format: textFormat, mode.value(for: input), style.rawValue, input),
suffix: suffix
)
}
}
Expand Down
2 changes: 1 addition & 1 deletion Sources/TranslatorSupport/TranslatorService.swift
Original file line number Diff line number Diff line change
Expand Up @@ -152,7 +152,7 @@ public final class TranslatorService {

translatedString = ""

for try await output in llamaModel.complete(prompt) {
for try await output in llamaModel.complete(prompt.text, suffix: prompt.suffix) {
translatedString.append(output)
}
})
Expand Down

0 comments on commit 3748999

Please sign in to comment.