@@ -136,6 +136,24 @@ class ChineseSpellingCorrector {
136136 return text
137137 }
138138
139+ // 运行推理,获取 logits 指针
140+ guard let ( logitsPtr, outputValue) = runInference ( inputIds: inputIds, seqLen: seqLen) else {
141+ return text
142+ }
143+ defer { api. pointee. ReleaseValue ( outputValue) }
144+
145+ // 应用纠正
146+ return applyCorrections (
147+ chars: chars,
148+ inputIds: inputIds,
149+ seqLen: seqLen,
150+ logitsPtr: logitsPtr,
151+ originalText: text
152+ )
153+ }
154+
155+ /// 运行 ONNX 推理,返回 logits 指针和输出 OrtValue(调用方负责释放)
156+ private func runInference( inputIds: [ Int32 ] , seqLen: Int ) -> ( UnsafeMutablePointer < Float > , OpaquePointer ) ? {
139157 // 准备输入张量
140158 var inputIdsCopy = inputIds. map { Int64 ( $0) }
141159 var attentionMask = [ Int64] ( repeating: 1 , count: seqLen)
@@ -147,15 +165,14 @@ class ChineseSpellingCorrector {
147165 let attMaskValue = createTensor ( & attentionMask, shape: & shape) ,
148166 let tokenTypeValue = createTensor ( & tokenTypeIds, shape: & shape) else {
149167 logger. error ( " 创建输入张量失败 " )
150- return text
168+ return nil
151169 }
152170 defer {
153171 api. pointee. ReleaseValue ( inputIdsValue)
154172 api. pointee. ReleaseValue ( attMaskValue)
155173 api. pointee. ReleaseValue ( tokenTypeValue)
156174 }
157175
158- // 运行推理
159176 let inputNames : [ UnsafePointer < CChar > ? ] = [
160177 " input_ids " ,
161178 " attention_mask " ,
@@ -184,14 +201,13 @@ class ChineseSpellingCorrector {
184201 let msg = Self . errorMessage ( from: api, status: runStatus)
185202 logger. error ( " Run 失败: \( msg, privacy: . public) " )
186203 api. pointee. ReleaseStatus ( runStatus)
187- return text
204+ return nil
188205 }
189206
190207 guard let outputValue = outputValues [ 0 ] else {
191208 logger. error ( " 输出值为 nil " )
192- return text
209+ return nil
193210 }
194- defer { api. pointee. ReleaseValue ( outputValue) }
195211
196212 // 获取 logits 数据指针
197213 var outputData : UnsafeMutableRawPointer ? = nil
@@ -200,15 +216,27 @@ class ChineseSpellingCorrector {
200216 let msg = Self . errorMessage ( from: api, status: dataStatus)
201217 logger. error ( " GetTensorMutableData 失败: \( msg, privacy: . public) " )
202218 api. pointee. ReleaseStatus ( dataStatus)
203- return text
219+ api. pointee. ReleaseValue ( outputValue)
220+ return nil
204221 }
205222
206223 guard let logitsPtr = outputData? . assumingMemoryBound ( to: Float . self) else {
207224 logger. error ( " logits 指针为 nil " )
208- return text
225+ api. pointee. ReleaseValue ( outputValue)
226+ return nil
209227 }
210228
211- // 对每个位置取 argmax,构建输出
229+ return ( logitsPtr, outputValue)
230+ }
231+
232+ /// 根据 logits 对每个位置进行纠正判定,返回纠正后的文本
233+ private func applyCorrections(
234+ chars: [ Character ] ,
235+ inputIds: [ Int32 ] ,
236+ seqLen: Int ,
237+ logitsPtr: UnsafeMutablePointer < Float > ,
238+ originalText: String
239+ ) -> String {
212240 let vocabSize = tokenizer. vocabSize
213241 var corrected = chars
214242 var correctionCount = 0
@@ -223,44 +251,14 @@ class ChineseSpellingCorrector {
223251 guard BertTokenizer . isChinese ( chars [ charIndex] ) else { continue }
224252 chineseCharCount += 1
225253
226- // 取 argmax
227- let logitsOffset = i * vocabSize
228- var maxVal : Float = - Float. infinity
229- var maxIdx : Int32 = 0
230-
231- for j in 0 ..< vocabSize {
232- let val = logitsPtr [ logitsOffset + j]
233- if val > maxVal {
234- maxVal = val
235- maxIdx = Int32 ( j)
236- }
237- }
238-
239- // 如果预测不同于输入且不是特殊 token
240- let originalId = inputIds [ i]
241- if maxIdx != originalId && maxIdx != tokenizer. unkId {
242- let originalLogit = logitsPtr [ logitsOffset + Int( originalId) ]
243- let logitDiff = maxVal - originalLogit
244-
245- // 条件 1:logit 差值必须足够大(提高到 5.0,防止低置信度替换)
246- guard logitDiff > 5.0 else { continue }
247-
248- // 条件 2:计算 softmax 概率,top-1 概率必须 > 0.9
249- // 为数值稳定性,对 logits 减去 maxVal 后再 softmax
250- var expSum : Float = 0
251- for j in 0 ..< vocabSize {
252- expSum += exp ( logitsPtr [ logitsOffset + j] - maxVal)
253- }
254- let topProb = 1.0 / expSum // exp(0) / expSum since maxVal - maxVal = 0
255- guard topProb > 0.9 else { continue }
256-
257- let correctedChar = tokenizer. decode ( [ maxIdx] )
258- if !correctedChar. isEmpty && correctedChar. count == 1 {
259- if let firstChar = correctedChar. first, BertTokenizer . isChinese ( firstChar) {
260- corrected [ charIndex] = firstChar
261- correctionCount += 1
262- }
263- }
254+ if let correctedChar = evaluateCorrection (
255+ tokenIndex: i,
256+ originalId: inputIds [ i] ,
257+ logitsPtr: logitsPtr,
258+ vocabSize: vocabSize
259+ ) {
260+ corrected [ charIndex] = correctedChar
261+ correctionCount += 1
264262 }
265263 }
266264
@@ -269,15 +267,63 @@ class ChineseSpellingCorrector {
269267 let correctionRatio = Float ( correctionCount) / Float( chineseCharCount)
270268 if correctionRatio > 0.2 {
271269 logger. warning ( " CSC 纠正比例过高 ( \( correctionCount) / \( chineseCharCount) = \( String ( format: " %.0f%% " , correctionRatio * 100 ) ) ),丢弃纠正结果 " )
272- return text
270+ return originalText
273271 }
274272
275273 let result = String ( corrected)
276- logger. info ( " CSC 纠正了 \( correctionCount) 处: \( text , privacy: . public) → \( result, privacy: . public) " )
274+ logger. info ( " CSC 纠正了 \( correctionCount) 处: \( originalText , privacy: . public) → \( result, privacy: . public) " )
277275 return result
278276 }
279277
280- return text
278+ return originalText
279+ }
280+
281+ /// 评估单个位置是否需要纠正,返回纠正后的字符(若不需要纠正则返回 nil)
282+ private func evaluateCorrection(
283+ tokenIndex: Int ,
284+ originalId: Int32 ,
285+ logitsPtr: UnsafeMutablePointer < Float > ,
286+ vocabSize: Int
287+ ) -> Character ? {
288+ let logitsOffset = tokenIndex * vocabSize
289+
290+ // 取 argmax
291+ var maxVal : Float = - Float. infinity
292+ var maxIdx : Int32 = 0
293+
294+ for j in 0 ..< vocabSize {
295+ let val = logitsPtr [ logitsOffset + j]
296+ if val > maxVal {
297+ maxVal = val
298+ maxIdx = Int32 ( j)
299+ }
300+ }
301+
302+ // 如果预测不同于输入且不是特殊 token
303+ guard maxIdx != originalId && maxIdx != tokenizer. unkId else { return nil }
304+
305+ let originalLogit = logitsPtr [ logitsOffset + Int( originalId) ]
306+ let logitDiff = maxVal - originalLogit
307+
308+ // 条件 1:logit 差值必须足够大(提高到 5.0,防止低置信度替换)
309+ guard logitDiff > 5.0 else { return nil }
310+
311+ // 条件 2:计算 softmax 概率,top-1 概率必须 > 0.9
312+ // 为数值稳定性,对 logits 减去 maxVal 后再 softmax
313+ var expSum : Float = 0
314+ for j in 0 ..< vocabSize {
315+ expSum += exp ( logitsPtr [ logitsOffset + j] - maxVal)
316+ }
317+ let topProb = 1.0 / expSum // exp(0) / expSum since maxVal - maxVal = 0
318+ guard topProb > 0.9 else { return nil }
319+
320+ let correctedChar = tokenizer. decode ( [ maxIdx] )
321+ if !correctedChar. isEmpty && correctedChar. count == 1 {
322+ if let firstChar = correctedChar. first, BertTokenizer . isChinese ( firstChar) {
323+ return firstChar
324+ }
325+ }
326+ return nil
281327 }
282328
283329 /// 创建 Int64 类型的 OrtValue 张量
0 commit comments