Named Entities & Navigation (#7)

codefiesta · web-flow · commit 6f7b86e1653b · 2025-05-12T14:02:27.000-07:00
diff --git a/Package.swift b/Package.swift
@@ -17,7 +17,7 @@ let package = Package(
         ),
     ],
     dependencies: [
-        .package(url: "https://github.com/codefiesta/VimKit", from: .init(0, 4, 7))
+        .package(url: "https://github.com/codefiesta/VimKit", from: .init(0, 4, 8))
     ],
     targets: [
         .target(
diff --git a/README.md b/README.md
@@ -24,32 +24,36 @@ Some examples of categorized actions include (but not limited to):
 ### Label Scheme
 The base model was created with the OntoNotes 5.0 NER annotations which includes:
 
-* **PERSON**: Individual names (e.g., Barack Obama).
-* **ORGANIZATION**: Company or institution names (e.g., Apple).
-* **LOCATION**: Geographical places (e.g., Tokyo).
+* **CARDINAL**: Cardinal numbers (e.g., 1, 2, 3).
 * **DATE**: Dates (e.g., May 8, 2025).
-* **TIME**: Times (e.g., 10:00 AM).
 * **EVENT**: Names of events (e.g., World Series).
-* **WORK\_OF\_ART**: Names of works of art (e.g., "Hamlet").
 * **FAC**: Buildings or facilities (e.g., White House).
 * **GPE**: Geo-political entities (e.g., United States).
 * **LANGUAGE**: Names of languages (e.g., English).
 * **LAW**: Legal names (e.g., The Constitution).
-* **NORP**: National/religious/political group (e.g., Democrats).
-* **CARDINAL**: Cardinal numbers (e.g., 1, 2, 3).
+* **LOC**: Represents locations (e.g., "New York City").
+* **MONEY**: Indicates monetary values (e.g., "100 dollars").
+* **NORP**: Represents national or political or religious groups (e.g., "Democrats", "the Catholic Church").
+* **ORDINAL**: Denotes ordinal numbers (e.g., "first", "second", "10th").
+* **ORG**: Represents organizations (e.g., "Google", "Microsoft").
+* **PERCENT**: Denotes percentages (e.g., "10%", "20%").
+* **PERSON**: Individual names (e.g., Barack Obama).
+* **PRODUCT**: Represents products (e.g., "iPhone", "MacBook").
+* **QUANTITY**: Indicates measurements or quantities (e.g., "10 kilograms").
+* **TIME**: Times (e.g., 10:00 AM).
+* **WORK\_OF\_ART**: Names of works of art (e.g., "Hamlet").
 
 The trained model provides Construction NER annotations:
 
-* **CON-BIM-CATG**: BIM Category - a high-level classification for families and elements, grouping them based on their functional type.
-* **CON-BIM-FAML**: BIM Family - a collection of elements that share common properties, behaviors, and physical characteristics.
-* **CON-BIM-TYPE**: BIM Type - a specific instantiation of a family that defines a unique set of parameters, essentially a variation within a family. Think of it as a specific size, material, or configuration of a particular family, such as a 3' x 6' door within a door family.
-* **CON-BIM-INST**: BIM Instance - a single, unique occurrence of a family type placed within a model.
-* **CON-BIM-LEVL**: BIM Level - a horizontal  plane used to define the vertical position of elements like walls, floors, and ceilings.
-* **CON-BIM-VIEW**: BIM View - represents a specific way of looking at the model, whether it's a 2D plan, elevation, section, or 3D view.
-
+* **CON\_BIM\_CATG**: BIM Category - a high-level classification for families and elements, grouping them based on their functional type.
+* **CON\_BIM\_FAML**: BIM Family - a collection of elements that share common properties, behaviors, and physical characteristics.
+* **CON\_BIM\_TYPE**: BIM Type - a specific instantiation of a family that defines a unique set of parameters, essentially a variation within a family. Think of it as a specific size, material, or configuration of a particular family, such as a 3' x 6' door within a door family.
+* **CON\_BIM\_INST**: BIM Instance - a single, unique occurrence of a family type placed within a model.
+* **CON\_BIM\_LEVL**: BIM Level - a horizontal  plane used to define the vertical position of elements like walls, floors, and ceilings.
+* **CON\_BIM\_VIEW**: BIM View - represents a specific way of looking at the model, whether it's a 2D plan, elevation, section, or 3D view.
 
 
-| Component    | Labels |
+| Component | Labels |
 | -------- | ------- |
-| named entities  | CARDINAL, DATE, EVENT, FAC, GPE, LANGUAGE, LAW, LOC, MONEY, NORP, ORDINAL, ORG, PERCENT, PERSON, PRODUCT, QUANTITY, TIME, WORK_OF_ART, CON-BIM-CATG, CON-BIM-FAML, CON-BIM-TYPE, CON-BIM-INST, CON-BIM-LEVL, CON-BIM-VIEW |
-| categories | ISOLATE, HIDE, QUANTIFY |
+| named entities  | CARDINAL, DATE, EVENT, FAC, GPE, LANGUAGE, LAW, LOC, MONEY, NORP, ORDINAL, ORG, PERCENT, PERSON, PRODUCT, QUANTITY, TIME, WORK\_OF\_ART, CON\_BIM\_CATG, CON\_BIM\_FAML, CON\_BIM\_TYPE, CON\_BIM\_INST, CON\_BIM\_LEVL, CON\_BIM\_VIEW |
+| categories | ISOLATE, HIDE, QUANTIFY, ZOOM\_IN, ZOOM\_OUT, PAN\_LEFT, PAN\_RIGHT, PAN\_UP, PAN\_DOWN, LOOK\_LEFT, LOOK\_RIGHT, LOOK\_UP, LOOK\_DOWN |
diff --git a/Sources/VimAssistant/Model/VimAssistant+Handler.swift b/Sources/VimAssistant/Model/VimAssistant+Handler.swift
@@ -19,52 +19,96 @@ public extension VimAssistant {
         ///   - prediction: the prediction
         func handle(vim: Vim, prediction: VimPrediction?) {
             guard let prediction, let bestPrediction = prediction.bestPrediction, bestPrediction.confidence >= 0.85 else { return }
-            guard prediction.entities.isNotEmpty else { return }
+            let action = bestPrediction.action
+            let ids = collect(vim: vim, prediction: prediction)
+            Task { @MainActor in
+                switch action {
+                case .hide:
+                    guard ids.isNotEmpty else { return }
+                    await vim.hide(ids: ids)
+                case .isolate:
+                    guard ids.isNotEmpty else { return }
+                    await vim.isolate(ids: ids)
+                case .quantify:
+                    // TODO: Probably just emit an event that shows the quantities view
+                    break
+                case .zoomIn:
+                    vim.zoom()
+                case .zoomOut:
+                    vim.zoom(out: true)
+                case .lookLeft:
+                    vim.look(.left)
+                case .lookRight:
+                    vim.look(.right)
+                case .lookUp:
+                    vim.look(.up)
+                case .lookDown:
+                    vim.look(.down)
+                case .panLeft:
+                    vim.pan(.left)
+                case .panRight:
+                    vim.pan(.right)
+                case .panUp:
+                    vim.pan(.up)
+                case .panDown:
+                    vim.pan(.down)
+                }
+            }
+        }
 
+        private func collect(vim: Vim, prediction: VimPrediction) -> [Int] {
+
+            guard let bestPrediction = prediction.bestPrediction, prediction.entities.isNotEmpty else { return []}
             let action = bestPrediction.action
 
-            print("❤️", bestPrediction)
+            switch action {
+            case .hide, .isolate:
+                guard let db = vim.db, db.nodes.isNotEmpty else { return [] }
+                let modelContext = ModelContext(db.modelContainer)
 
-            for entity in prediction.entities {
-                if entity.label == "CON-BIM-CATG" {
-                    print("🚀", entity.value)
-                    perform(vim: vim, action: action, category: entity.value)
-                } else if entity.label == "CON-BIM-FAML" {
+                var ids: Set<Int> = .init()
 
-                } else if entity.label == "CON-BIM-TYPE" {
+                // Fetch all geometry nodes
+                let nodes = db.nodes
+                let predicate = Database.Node.predicate(nodes: nodes)
+                let descriptor = FetchDescriptor<Database.Node>(predicate: predicate, sortBy: [SortDescriptor(\.index)])
+                guard let results = try? modelContext.fetch(descriptor), results.isNotEmpty else { return [] }
 
-                }
-            }
-        }
+                let categoryNames = prediction.entities.filter{ $0.label == .bimCategory }.map { $0.value }
+                let familyNames = prediction.entities.filter{ $0.label == .bimFamily }.map { $0.value }
 
-        private func perform(vim: Vim, action: VimPrediction.Action, category: String) {
+                // Tuple of category names and ids
+                let categories = results.compactMap{ $0.element?.category?.name }.uniqued().sorted{ $0 < $1 }.map { name in
+                    (name: name, ids: results.filter{ $0.element?.category?.name == name}.compactMap{ Int($0.index) })
+                }
 
-            guard let db = vim.db else { return }
-            let modelContext = ModelContext(db.modelContainer)
+                // Tuple of family names and ids
+                let familes = results.compactMap{ $0.element?.familyName }.uniqued().sorted{ $0 < $1 }.map { name in
+                    (name: name, ids: results.filter{ $0.element?.familyName == name}.compactMap{ Int($0.index) })
+                }
 
-            let orderedSame = ComparisonResult.orderedSame
-            let predicate = #Predicate<Database.Node>{
-                if let element = $0.element, let cat = element.category {
-                    return cat.name.caseInsensitiveCompare(category) == orderedSame
-                } else {
-                    return false
+                // Collect the ids of the matching categories
+                for name in categoryNames {
+                    let found = categories.filter{ name.localizedStandardContains($0.name) }.map{ $0.ids }.reduce([], +)
+                    ids.formUnion(found)
                 }
-            }
 
-            let descriptor = FetchDescriptor<Database.Node>(predicate: predicate, sortBy: [SortDescriptor(\.index)])
-            guard let results = try? modelContext.fetch(descriptor), results.isNotEmpty else { return }
-            let ids = results.compactMap{ Int($0.index) }
-            Task {
-                switch action {
-                case .hide:
-                    await vim.hide(ids: ids)
-                case .isolate:
-                    await vim.isolate(ids: ids)
-                case .quantify:
-                    break
+                // Collect the ids of the matching families
+                for name in familyNames {
+                    let found = familes.filter{ name.localizedStandardContains($0.name) }.map{ $0.ids }.reduce([], +)
+                    ids.formUnion(found)
                 }
+                return ids.sorted()
+            case .quantify, .zoomIn, .zoomOut, .lookLeft, .lookRight, .lookUp, .lookDown, .panLeft, .panRight, .panUp, .panDown:
+                return []
             }
+
         }
+    }
+}
 
+extension Array where Element == String {
+    func containsIgnoringCase(_ element: Element) -> Bool {
+        contains { $0.caseInsensitiveCompare(element) == .orderedSame }
     }
 }
diff --git a/Sources/VimAssistant/Model/VimPrediction.swift b/Sources/VimAssistant/Model/VimPrediction.swift
@@ -16,10 +16,48 @@ public struct VimPrediction: Decodable, Equatable {
         case tokens = "tokens"
     }
 
+    enum NerLabel: String, Identifiable {
+
+        case person = "PERSON"
+        case organization = "ORGANIZATION"
+        case location = "LOCATION"
+        case date = "DATE"
+        case time = "TIME"
+        case event = "EVENT"
+        case workOfArt = "WORK_OF_ART"
+        case fac = "FAC"
+        case gpe = "GPE"
+        case language = "LANGUAGE"
+        case law = "LAW"
+        case norp = "NORP"
+        case product = "PRODUCT"
+        case cardinal = "CARDINAL"
+        case bimCategory = "CON_BIM_CATG"
+        case bimFamily = "CON_BIM_FAML"
+        case bimType = "CON_BIM_TYPE"
+        case bimInstance = "CON_BIM_INST"
+        case bimLevel = "CON_BIM_LEVL"
+        case bimView = "CON_BIM_VIEW"
+
+        public var id: String {
+            rawValue
+        }
+    }
+
     enum Action: String, Codable, Identifiable {
         case isolate = "ISOLATE"
         case hide = "HIDE"
         case quantify = "QUANTIFY"
+        case zoomIn = "ZOOM_IN"
+        case zoomOut = "ZOOM_OUT"
+        case lookLeft = "LOOK_LEFT"
+        case lookRight = "LOOK_RIGHT"
+        case lookUp = "LOOK_UP"
+        case lookDown = "LOOK_DOWN"
+        case panLeft = "PAN_LEFT"
+        case panRight = "PAN_RIGHT"
+        case panUp = "PAN_UP"
+        case panDown = "PAN_DOWN"
 
         public var id: String {
             rawValue
@@ -76,17 +114,18 @@ public struct VimPrediction: Decodable, Equatable {
             case end = "end"
         }
 
-        var label: String
+        var label: NerLabel
         var value: String = .empty
         var range: Range<Int>
 
         public var id: String {
-            label + "_\(range)"
+            label.rawValue + "_\(range)"
         }
 
         init(from decoder: Decoder) throws {
             let values = try decoder.container(keyedBy: CodingKeys.self)
-            label = try values.decode(String.self, forKey: .label)
+            let labelString = try values.decode(String.self, forKey: .label)
+            label = .init(rawValue: labelString)!
             let start = try values.decode(Int.self, forKey: .start)
             let end = try values.decode(Int.self, forKey: .end)
             range = start..<end
diff --git a/Sources/VimAssistant/Views/VimPredictionView.swift b/Sources/VimAssistant/Views/VimPredictionView.swift
@@ -71,7 +71,7 @@ struct VimPredictionView: View {
         for entity in prediction.entities {
             let entityText = text[entity.range]
             var attributedEntityString = AttributedString(entityText)
-            attributedEntityString.foregroundColor = .orange
+            attributedEntityString.foregroundColor = .cyan
             attributedEntityString.underlineStyle = .single
             attributedEntityString.link = URL(string: "/\(entity.label)/\(entity.value)")!
             result.replaceSubrange(bounds: entity.range, with: attributedEntityString)
@@ -101,10 +101,11 @@ struct VimPredictionView: View {
                 HStack {
                     Text(text[entity.range])
                         .bold()
-                    Text(entity.label)
-                        .padding(2)
-                        .background(Color.orange)
-                        .cornerRadius(4)
+                    Text(entity.label.rawValue)
+                        .padding(1)
+                        .background(Color.cyan)
+                        .foregroundStyle(Color.black)
+                        .cornerRadius(2)
                 }
             }
             if let bestPrediction = prediction.bestPrediction {
@@ -116,7 +117,7 @@ struct VimPredictionView: View {
                 HStack {
                     Text(bestPrediction.action.rawValue.lowercased())
                         .bold()
-                    Text(bestPrediction.confidence.formatted(.percent))
+                    Text(bestPrediction.confidence.formatted(.percent.precision(.fractionLength(2))))
                         .foregroundStyle(predictionConfidenceColor)
                 }
             }
@@ -128,7 +129,7 @@ struct VimPredictionView: View {
 
 #Preview {
 
-    let json = "{\"text\":\"Hide all walls and air terminals \",\"ents\":[{\"start\":9,\"end\":14,\"label\":\"CON-BIM-CATG\"},{\"start\":19,\"end\":32,\"label\":\"CON-BIM-CATG\"}],\"cats\":{\"ISOLATE\":0.0122569752857089,\"HIDE\":0.978784739971161,\"QUANTIFY\":0.00895828753709793}}"
+    let json = "{\"text\":\"Hide all walls and air terminals \",\"ents\":[{\"start\":9,\"end\":14,\"label\":\"CON_BIM_CATG\"},{\"start\":19,\"end\":32,\"label\":\"CON_BIM_CATG\"}],\"cats\":{\"ISOLATE\":0.0122569752857089,\"HIDE\":0.978784739971161,\"QUANTIFY\":0.00895828753709793}}"
     let prediction = try! JSONDecoder().decode(VimPrediction.self, from: json.data(using: .utf8)!)
-    VimPredictionView(prediction: prediction, explain: .constant(false))
+    VimPredictionView(prediction: prediction, explain: .constant(true))
 }
