Named Entity Extraction - Swift

ner.swift

func getEntities (_ sentence : String) -> [String] {
    let tagger = NSLinguisticTagger(tagSchemes: [.nameType], options: 0)
    let options: NSLinguisticTagger.Options = [.omitPunctuation, .omitWhitespace]
    
    tagger.string = sentence
    let range = NSRange(location:0, length: sentence.utf16.count)
    
    var prevTag : NSLinguisticTag? = nil
    var accName : String = ""
    
    let selectedTags : [NSLinguisticTag] = [.personalName, .organizationName, .placeName]
    
    var entities : [String] = []
    
    tagger.enumerateTags(in: range, unit: .word, scheme: .nameType, options: options) { tag, tokenRange, stop in
        if let tag = tag {
            let name = (sentence as NSString).substring(with: tokenRange)
            if let prevTag = prevTag, prevTag == tag {
                accName += " " + name
            } else {
                if prevTag != nil && selectedTags.contains(prevTag!) {
                    entities.append(accName)
                }
                accName = name
            }
            prevTag = tag
        }
    }
    if prevTag != nil && selectedTags.contains(prevTag!) {
        entities.append(accName)
    }
    
    return entities
}