Add Sequence and Bert processors (#129)

Author: DePasqualeOrg

Sources/Tokenizers/PostProcessor.swift

@@ -25,6 +25,8 @@ enum PostProcessorType: String {
     case TemplateProcessing
     case ByteLevel
     case RobertaProcessing
+    case BertProcessing
+    case Sequence
 }
 
 struct PostProcessorFactory {
@@ -33,10 +35,12 @@ struct PostProcessorFactory {
         guard let typeName = config.type?.stringValue else { return nil }
         let type = PostProcessorType(rawValue: typeName)
         switch type {
-        case .TemplateProcessing: return TemplateProcessing(config: config)
-        case .ByteLevel         : return ByteLevelPostProcessor(config: config)
-        case .RobertaProcessing : return RobertaProcessing(config: config)
-        default                 : fatalError("Unsupported PostProcessor type: \(typeName)")
+            case .TemplateProcessing : return TemplateProcessing(config: config)
+            case .ByteLevel          : return ByteLevelPostProcessor(config: config)
+            case .RobertaProcessing  : return RobertaProcessing(config: config)
+            case .BertProcessing     : return BertProcessing(config: config)
+            case .Sequence           : return SequenceProcessing(config: config)
+            default                  : fatalError("Unsupported PostProcessor type: \(typeName)")
         }
     }
 }
@@ -139,3 +143,51 @@ class RobertaProcessing: PostProcessor {
         return text.reversed().prefix(while: { $0.isWhitespace }).count - 1
     }
 }
+
+class BertProcessing: PostProcessor {
+    private let sep: (UInt, String)
+    private let cls: (UInt, String)
+
+    required public init(config: Config) {
+        guard let sep = config.sep?.tokenValue else { fatalError("Missing `sep` processor configuration") }
+        guard let cls = config.cls?.tokenValue else { fatalError("Missing `cls` processor configuration") }
+        self.sep = sep
+        self.cls = cls
+    }
+
+    func postProcess(tokens: [String], tokensPair: [String]?, addSpecialTokens: Bool = true) -> [String] {
+        guard addSpecialTokens else { return tokens + (tokensPair ?? []) }
+
+        var outTokens = [self.cls.1] + tokens + [self.sep.1]
+        if let tokensPair = tokensPair, !tokensPair.isEmpty {
+            outTokens += tokensPair + [self.sep.1]
+        }
+
+        return outTokens
+    }
+}
+
+class SequenceProcessing: PostProcessor {
+    private let processors: [PostProcessor]
+
+    required public init(config: Config) {
+        guard let processorConfigs = config.processors?.arrayValue else {
+            fatalError("Missing `processors` configuration")
+        }
+
+        self.processors = processorConfigs.compactMap { PostProcessorFactory.fromConfig(config: $0) }
+    }
+
+    func postProcess(tokens: [String], tokensPair: [String]?, addSpecialTokens: Bool = true) -> [String] {
+        var currentTokens = tokens
+        var currentTokensPair = tokensPair
+
+        for processor in processors {
+            let processed = processor.postProcess(tokens: currentTokens, tokensPair: currentTokensPair, addSpecialTokens: addSpecialTokens)
+            currentTokens = processed
+            currentTokensPair = nil  // After the first processor, we no longer have a separate pair
+        }
+
+        return currentTokens
+    }
+}