diff --git a/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift b/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift index ca6ebaa..4f8999f 100644 --- a/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift +++ b/Sources/SimilaritySearchKit/Core/Embeddings/Metrics/DistanceMetrics.swift @@ -6,6 +6,7 @@ // import Foundation +import Accelerate /// A struct implementing the `DistanceMetricProtocol` using the dot product. /// @@ -20,10 +21,22 @@ public struct DotProduct: DistanceMetricProtocol { return sortedScores(scores: scores, topK: resultsCount) } - public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { - let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +) - return dotProduct - } +// public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { +// let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +) +// return dotProduct +// } + + public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { + // Ensure the embeddings have the same length + precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length") + + var dotProduct: Float = 0 + + // Calculate dot product using Accelerate + vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count)) + + return dotProduct + } } @@ -40,14 +53,36 @@ public struct CosineSimilarity: DistanceMetricProtocol { return sortedScores(scores: scores, topK: resultsCount) } - public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { - // Calculate cosine distance - let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +) - let firstMagnitude = sqrt(firstEmbedding.map { $0 * $0 }.reduce(0, +)) - let secondMagnitude = sqrt(secondEmbedding.map { $0 * $0 }.reduce(0, +)) - - return dotProduct / (firstMagnitude * secondMagnitude) - } +// public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { +// // Calculate cosine distance +// let dotProduct = zip(firstEmbedding, secondEmbedding).map(*).reduce(0, +) +// let firstMagnitude = sqrt(firstEmbedding.map { $0 * $0 }.reduce(0, +)) +// let secondMagnitude = sqrt(secondEmbedding.map { $0 * $0 }.reduce(0, +)) +// +// return dotProduct / (firstMagnitude * secondMagnitude) +// } + + public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { + // Ensure the embeddings have the same length + precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length") + + var dotProduct: Float = 0 + var firstMagnitude: Float = 0 + var secondMagnitude: Float = 0 + + // Calculate dot product and magnitudes using Accelerate + vDSP_dotpr(firstEmbedding, 1, secondEmbedding, 1, &dotProduct, vDSP_Length(firstEmbedding.count)) + vDSP_svesq(firstEmbedding, 1, &firstMagnitude, vDSP_Length(firstEmbedding.count)) + vDSP_svesq(secondEmbedding, 1, &secondMagnitude, vDSP_Length(secondEmbedding.count)) + + // Take square root of magnitudes + firstMagnitude = sqrt(firstMagnitude) + secondMagnitude = sqrt(secondMagnitude) + + // Return cosine similarity + return dotProduct / (firstMagnitude * secondMagnitude) + } + } /// A struct implementing the `DistanceMetricProtocol` using Euclidean distance. @@ -64,10 +99,23 @@ public struct EuclideanDistance: DistanceMetricProtocol { return sortedDistances(distances: distances, topK: resultsCount) } - public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { - let squaredDifferences = zip(firstEmbedding, secondEmbedding).map { ($0 - $1) * ($0 - $1) } - return sqrt(squaredDifferences.reduce(0, +)) - } +// public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { +// let squaredDifferences = zip(firstEmbedding, secondEmbedding).map { ($0 - $1) * ($0 - $1) } +// return sqrt(squaredDifferences.reduce(0, +)) +// } + + public func distance(between firstEmbedding: [Float], and secondEmbedding: [Float]) -> Float { + // Ensure the embeddings have the same length + precondition(firstEmbedding.count == secondEmbedding.count, "Embeddings must have the same length") + + var distance: Float = 0 + + // Calculate squared differences and sum them using Accelerate + vDSP_distancesq(firstEmbedding, 1, secondEmbedding, 1, &distance, vDSP_Length(firstEmbedding.count)) + + // Return the square root of the summed squared differences + return sqrt(distance) + } } // MARK: - Helpers diff --git a/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift b/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift index 643d3bf..04f22ad 100644 --- a/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift +++ b/Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift @@ -319,6 +319,7 @@ extension SimilarityIndex { @available(macOS 13.0, iOS 16.0, *) extension SimilarityIndex { + public func saveIndex(toDirectory path: URL? = nil, name: String? = nil) throws -> URL { let indexName = name ?? self.indexName let basePath: URL