Skip to content

Commit 9bec547

Browse files
ZachNagengastZachNagengast
committed
Formatting, and replace precondition with default value to prevent crashing
1 parent 6f02580 commit 9bec547

3 files changed

Lines changed: 83 additions & 85 deletions

File tree

‎Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift‎

Lines changed: 55 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -5,8 +5,8 @@
55
// Created by Zach Nagengast on 4/7/23.
66
//
77

8-
import Foundation
98
import Accelerate
9+
import Foundation
1010

1111
/// A struct implementing the `DistanceMetricProtocol` using the dot product.
1212
///
@@ -21,19 +21,20 @@ public struct DotProduct: DistanceMetricProtocol {
2121
return sortedScores(scores: scores, topK: resultsCount)
2222
}
2323

24-
25-
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
26-
// Ensure the embeddings have the same length
27-
precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
28-
29-
var dotProduct: Float = 0
30-
31-
// Calculate dot product using Accelerate
32-
vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
33-
34-
return dotProduct
35-
}
36-
24+
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
25+
// Ensure the embeddings have the same length
26+
if firstEmbedding.count != secondEmbedding.count {
27+
print("Embeddings must have the same length")
28+
return -Float.greatestFiniteMagnitude
29+
}
30+
31+
var dotProduct: Float = 0
32+
33+
// Calculate dot product using Accelerate
34+
vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
35+
36+
return dotProduct
37+
}
3738
}
3839

3940
/// A struct implementing the `DistanceMetricProtocol` using cosine similarity.
@@ -49,28 +50,29 @@ public struct CosineSimilarity: DistanceMetricProtocol {
4950
return sortedScores(scores: scores, topK: resultsCount)
5051
}
5152

52-
53-
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
54-
// Ensure the embeddings have the same length
55-
precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
56-
57-
var dotProduct: Float = 0
58-
var firstMagnitude: Float = 0
59-
var secondMagnitude: Float = 0
60-
61-
// Calculate dot product and magnitudes using Accelerate
62-
vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
63-
vDSP_svesq(firstEmbedding, 1, &firstMagnitude, vDSP_Length(firstEmbedding.count))
64-
vDSP_svesq(secondEmbedding, 1, &secondMagnitude, vDSP_Length(secondEmbedding.count))
65-
66-
// Take square root of magnitudes
67-
firstMagnitude = sqrt(firstMagnitude)
68-
secondMagnitude = sqrt(secondMagnitude)
69-
70-
// Return cosine similarity
71-
return dotProduct / (firstMagnitude * secondMagnitude)
72-
}
73-
53+
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
54+
// Ensure the embeddings have the same length
55+
if firstEmbedding.count != secondEmbedding.count {
56+
print("Embeddings must have the same length")
57+
return -1
58+
}
59+
60+
var dotProduct: Float = 0
61+
var firstMagnitude: Float = 0
62+
var secondMagnitude: Float = 0
63+
64+
// Calculate dot product and magnitudes using Accelerate
65+
vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count))
66+
vDSP_svesq(firstEmbedding, 1, &firstMagnitude, vDSP_Length(firstEmbedding.count))
67+
vDSP_svesq(secondEmbedding, 1, &secondMagnitude, vDSP_Length(secondEmbedding.count))
68+
69+
// Take square root of magnitudes
70+
firstMagnitude = sqrt(firstMagnitude)
71+
secondMagnitude = sqrt(secondMagnitude)
72+
73+
// Return cosine similarity
74+
return dotProduct / (firstMagnitude * secondMagnitude)
75+
}
7476
}
7577

7678
/// A struct implementing the `DistanceMetricProtocol` using Euclidean distance.
@@ -79,26 +81,28 @@ public struct CosineSimilarity: DistanceMetricProtocol {
7981
///
8082
/// - Note: Use this metric when the magnitudes of the embeddings are significant in your use case, and the embeddings are distributed in a Euclidean space.
8183
public struct EuclideanDistance: DistanceMetricProtocol {
82-
8384
public init() {}
8485

8586
public func findNearest(for queryEmbedding: [Float], in neighborEmbeddings: [[Float]], resultsCount: Int) -> [(Float, Int)] {
8687
let distances = neighborEmbeddings.map { distance(between: queryEmbedding, and: $0) }
8788
return sortedDistances(distances: distances, topK: resultsCount)
8889
}
89-
90-
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
91-
// Ensure the embeddings have the same length
92-
precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length")
93-
94-
var distance: Float = 0
95-
96-
// Calculate squared differences and sum them using Accelerate
97-
vDSP_distancesq(firstEmbedding, 1, secondEmbedding, 1, &distance, vDSP_Length(firstEmbedding.count))
98-
99-
// Return the square root of the summed squared differences
100-
return sqrt(distance)
101-
}
90+
91+
public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float {
92+
// Ensure the embeddings have the same length
93+
if firstEmbedding.count != secondEmbedding.count {
94+
print("Embeddings must have the same length")
95+
return Float.greatestFiniteMagnitude
96+
}
97+
98+
var distance: Float = 0
99+
100+
// Calculate squared differences and sum them using Accelerate
101+
vDSP_distancesq(firstEmbedding, 1, secondEmbedding, 1, &distance, vDSP_Length(firstEmbedding.count))
102+
103+
// Return the square root of the summed squared differences
104+
return sqrt(distance)
105+
}
102106
}
103107

104108
// MARK: - Helpers
@@ -133,7 +137,7 @@ public func sortedScores(scores: [Float], topK: Int) -> [(Float, Int)] {
133137
/// - Parameters:
134138
/// - distances: An array of Float values representing distances.
135139
/// - topK: The number of top distances to return.
136-
///
140+
///
137141
/// - Returns: An array of tuples containing the top K distances and their corresponding indices.
138142
public func sortedDistances(distances: [Float], topK: Int) -> [(Float, Int)] {
139143
// Combine indices & distances

‎Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift‎

Lines changed: 28 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@
66
//
77

88
import Foundation
9-
import OSLog
109

1110
// MARK: - Type Aliases
1211

@@ -19,20 +18,18 @@ public typealias VectorStoreType = SimilarityIndex.VectorStoreType
1918

2019
@available(macOS 11.0, iOS 15.0, *)
2120
public class SimilarityIndex: Identifiable, Hashable {
22-
23-
public static func == (lhs: SimilarityIndex, rhs: SimilarityIndex) -> Bool {
24-
return lhs.id == rhs.id
25-
}
26-
27-
public func hash(into hasher: inout Hasher) {
28-
hasher.combine(id)
29-
}
30-
3121
// MARK: - Properties
3222

33-
/// A unique identifier
34-
public var id: UUID = UUID()
35-
23+
/// Unique identifier for this index instance
24+
public var id: UUID = .init()
25+
public static func == (lhs: SimilarityIndex, rhs: SimilarityIndex) -> Bool {
26+
return lhs.id == rhs.id
27+
}
28+
29+
public func hash(into hasher: inout Hasher) {
30+
hasher.combine(id)
31+
}
32+
3633
/// The items stored in the index.
3734
public var indexItems: [IndexItem] = []
3835

@@ -161,7 +158,7 @@ public class SimilarityIndex: Identifiable, Hashable {
161158
var indexIds: [String] = []
162159
var indexEmbeddings: [[Float]] = []
163160

164-
indexItems.forEach { item in
161+
for item in indexItems {
165162
indexIds.append(item.id)
166163
indexEmbeddings.append(item.embedding)
167164
}
@@ -220,18 +217,18 @@ public class SimilarityIndex: Identifiable, Hashable {
220217
// MARK: - CRUD
221218

222219
@available(macOS 11.0, iOS 15.0, *)
223-
extension SimilarityIndex {
220+
public extension SimilarityIndex {
224221
// MARK: Create
225222

226-
// Add an item with optional pre-computed embedding
227-
public func addItem(id: String, text: String, metadata: [String: String], embedding: [Float]? = nil) async {
223+
/// Add an item with optional pre-computed embedding
224+
func addItem(id: String, text: String, metadata: [String: String], embedding: [Float]? = nil) async {
228225
let embeddingResult = await getEmbedding(for: text, embedding: embedding)
229226

230227
let item = IndexItem(id: id, text: text, embedding: embeddingResult, metadata: metadata)
231228
indexItems.append(item)
232229
}
233230

234-
public func addItems(ids: [String], texts: [String], metadata: [[String: String]], embeddings: [[Float]?]? = nil, onProgress: ((String) -> Void)? = nil) async {
231+
func addItems(ids: [String], texts: [String], metadata: [[String: String]], embeddings: [[Float]?]? = nil, onProgress: ((String) -> Void)? = nil) async {
235232
// Check if all input arrays have the same length
236233
guard ids.count == texts.count, texts.count == metadata.count else {
237234
fatalError("Input arrays must have the same length.")
@@ -258,7 +255,7 @@ extension SimilarityIndex {
258255
}
259256
}
260257

261-
public func addItems(_ items: [IndexItem], completion: (() -> Void)? = nil) {
258+
func addItems(_ items: [IndexItem], completion: (() -> Void)? = nil) {
262259
Task {
263260
for item in items {
264261
await self.addItem(id: item.id, text: item.text, metadata: item.metadata, embedding: item.embedding)
@@ -269,17 +266,17 @@ extension SimilarityIndex {
269266

270267
// MARK: Read
271268

272-
public func getItem(id: String) -> IndexItem? {
269+
func getItem(id: String) -> IndexItem? {
273270
return indexItems.first { $0.id == id }
274271
}
275272

276-
public func sample(_ count: Int) -> [IndexItem]? {
273+
func sample(_ count: Int) -> [IndexItem]? {
277274
return Array(indexItems.prefix(upTo: count))
278275
}
279276

280277
// MARK: Update
281278

282-
public func updateItem(id: String, text: String? = nil, embedding: [Float]? = nil, metadata: [String: String]? = nil) {
279+
func updateItem(id: String, text: String? = nil, embedding: [Float]? = nil, metadata: [String: String]? = nil) {
283280
// Check if the provided embedding has the correct dimension
284281
if let embedding = embedding, embedding.count != dimension {
285282
print("Dimension mismatch, expected \(dimension), saw \(embedding.count)")
@@ -306,21 +303,20 @@ extension SimilarityIndex {
306303

307304
// MARK: Delete
308305

309-
public func removeItem(id: String) {
306+
func removeItem(id: String) {
310307
indexItems.removeAll { $0.id == id }
311308
}
312309

313-
public func removeAll() {
310+
func removeAll() {
314311
indexItems.removeAll()
315312
}
316313
}
317314

318315
// MARK: - Persistence
319316

320317
@available(macOS 13.0, iOS 16.0, *)
321-
extension SimilarityIndex {
322-
323-
public func saveIndex(toDirectory path: URL? = nil, name: String? = nil) throws -> URL {
318+
public extension SimilarityIndex {
319+
func saveIndex(toDirectory path: URL? = nil, name: String? = nil) throws -> URL {
324320
let indexName = name ?? self.indexName
325321
let basePath: URL
326322

@@ -333,13 +329,12 @@ extension SimilarityIndex {
333329

334330
let savedVectorStore = try vectorStore.saveIndex(items: indexItems, to: basePath, as: indexName)
335331

336-
let bundleId: String = Bundle.main.bundleIdentifier ?? "com.similarity-search-kit.logger"
337-
let logger: Logger = Logger(subsystem: bundleId, category: "similarityIndexSave")
332+
print("Saved \(indexItems.count) index items to \(savedVectorStore.absoluteString)")
338333

339334
return savedVectorStore
340335
}
341336

342-
public func loadIndex(fromDirectory path: URL? = nil, name: String? = nil) throws -> [IndexItem]? {
337+
func loadIndex(fromDirectory path: URL? = nil, name: String? = nil) throws -> [IndexItem]? {
343338
if let indexPath = try getIndexPath(fromDirectory: path, name: name) {
344339
indexItems = try vectorStore.loadIndex(from: indexPath)
345340
return indexItems
@@ -355,7 +350,7 @@ extension SimilarityIndex {
355350
/// - name: optional name
356351
///
357352
/// - Returns: an optional URL
358-
public func getIndexPath(fromDirectory path: URL? = nil, name: String? = nil) throws -> URL? {
353+
func getIndexPath(fromDirectory path: URL? = nil, name: String? = nil) throws -> URL? {
359354
let indexName = name ?? self.indexName
360355
let basePath: URL
361356

@@ -382,7 +377,7 @@ extension SimilarityIndex {
382377
return appSpecificDirectory
383378
}
384379

385-
public func estimatedSizeInBytes() -> Int {
380+
func estimatedSizeInBytes() -> Int {
386381
var totalSize = 0
387382

388383
for item in indexItems {
@@ -397,7 +392,7 @@ extension SimilarityIndex {
397392
let embeddingSize = item.embedding.count * floatSize
398393

399394
// Calculate the size of 'metadata' property
400-
let metadataSize = item.metadata.reduce(0) { (size, keyValue) -> Int in
395+
let metadataSize = item.metadata.reduce(0) { size, keyValue -> Int in
401396
let keySize = keyValue.key.utf8.count
402397
let valueSize = keyValue.value.utf8.count
403398
return size + keySize + valueSize

‎Sources/SimilaritySearchKit/Core/Persistence/Json/JsonStore.swift‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
import Foundation
99

1010
public class JsonStore: VectorStoreProtocol {
11-
1211
public func saveIndex(items: [IndexItem], to url: URL, as name: String) throws -> URL {
1312
let encoder = JSONEncoder()
1413
let data = try encoder.encode(items)

0 commit comments

Comments
 (0)