Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Refactor SimilarityIndex initialization sync and setupDimension method public #35

Closed
wants to merge 3 commits into from
Closed

Refactor SimilarityIndex initialization sync and setupDimension method public #35

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
7bf8423
"Refactor SimilarityIndex initialization and setupDimension method fo…
buhe Feb 13, 2024
b00ebd8
"Refactor loadIndex method to remove unnecessary closure and improve …
buhe Feb 14, 2024
39bbd0b
"Refactor loadIndex method to assign loaded index items directly to i…
buhe Feb 14, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 6 additions & 10 deletions Sources/SimilaritySearchKit/Core/Index/SimilarityIndex.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -101,18 +101,16 @@ public class SimilarityIndex {

// MARK: - Initializers

public init(name: String? = nil, model: (any EmbeddingsProtocol)? = nil, metric: (any DistanceMetricProtocol)? = nil, vectorStore: (any VectorStoreProtocol)? = nil) async {
public init(name: String? = nil, model: (any EmbeddingsProtocol)? = nil, metric: (any DistanceMetricProtocol)? = nil, vectorStore: (any VectorStoreProtocol)? = nil) {
// Setup index with defaults
self.indexName = name ?? "SimilaritySearchKitIndex"
self.indexModel = model ?? NativeEmbeddings()
self.indexMetric = metric ?? CosineSimilarity()
self.vectorStore = vectorStore ?? JsonStore()

// Run the model once to discover dimention size
await setupDimension()
}

private func setupDimension() async {
// Run the model once to discover dimention size
public func setupDimension() async {
if let testVector = await indexModel.encode(sentence: "Test sentence") {
dimension = testVector.count
} else {
Expand Down Expand Up @@ -326,11 +324,9 @@ extension SimilarityIndex {

public func loadIndex(fromDirectory path: URL? = nil, name: String? = nil) throws -> [IndexItem]? {
if let indexPath = try getIndexPath(fromDirectory: path, name: name) {
let loadedIndexItems = try vectorStore.loadIndex(from: indexPath)
addItems(loadedIndexItems) {[self] in
print("Loaded \(indexItems.count) index items from \(indexPath.absoluteString)")
}
return loadedIndexItems
indexItems = try vectorStore.loadIndex(from: indexPath)
print("Loaded \(indexItems.count) index items from \(indexPath.absoluteString)")
return indexItems
}

return nil
Expand Down
7 changes: 4 additions & 3 deletions Tests/SimilaritySearchKitTests/BenchmarkTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -41,7 +41,8 @@ class BenchmarkTests: XCTestCase {
let expectation = XCTestExpectation(description: "Encoding passage texts")

Task {
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings())
let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings())
await similarityIndex.setupDimension()
await similarityIndex.addItems(
ids: [UUID().uuidString],
texts: [searchPassage.text],
Expand Down Expand Up @@ -125,8 +126,8 @@ class BenchmarkTests: XCTestCase {

Task {
print("\nGenerating similarity index for \(testAmount) passages")
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings())

let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings())
await similarityIndex.setupDimension()
var startTime = CFAbsoluteTimeGetCurrent()
await similarityIndex.addItems(
ids: passageIds,
Expand Down
24 changes: 12 additions & 12 deletions Tests/SimilaritySearchKitTests/SimilaritySearchKitTests.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -21,8 +21,8 @@ class SimilaritySearchKitTests: XCTestCase {
}

func testSavingJsonIndex() async {
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())

let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())
await similarityIndex.setupDimension()
await similarityIndex.addItem(id: "1", text: "Example text", metadata: ["source": "test source"], embedding: [0.1, 0.2, 0.3])

let successPath = try! similarityIndex.saveIndex(name: "TestIndexForSaving")
Expand All @@ -31,24 +31,24 @@ class SimilaritySearchKitTests: XCTestCase {
}

func testLoadingJsonIndex() async {
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())

let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())
await similarityIndex.setupDimension()
await similarityIndex.addItem(id: "1", text: "Example text", metadata: ["source": "test source"])

let successPath = try! similarityIndex.saveIndex(name: "TestIndexForLoading")

XCTAssertNotNil(successPath)

let similarityIndex2 = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())

let similarityIndex2 = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: JsonStore())
await similarityIndex2.setupDimension()
let loadedItems = try! similarityIndex2.loadIndex(name: "TestIndexForLoading")

XCTAssertNotNil(loadedItems)
}

func testSavingBinaryIndex() async {
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())

let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())
await similarityIndex.setupDimension()
await similarityIndex.addItem(id: "1", text: "Example text", metadata: ["source": "test source"], embedding: [0.1, 0.2, 0.3])

let successPath = try! similarityIndex.saveIndex(name: "TestIndexForSaving")
Expand All @@ -57,16 +57,16 @@ class SimilaritySearchKitTests: XCTestCase {
}

func testLoadingBinaryIndex() async {
let similarityIndex = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())

let similarityIndex = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())
await similarityIndex.setupDimension()
await similarityIndex.addItem(id: "1", text: "Example text", metadata: ["source": "test source"])

let successPath = try! similarityIndex.saveIndex(name: "TestIndexForLoading")

XCTAssertNotNil(successPath)

let similarityIndex2 = await SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())

let similarityIndex2 = SimilarityIndex(model: DistilbertEmbeddings(), vectorStore: BinaryStore())
await similarityIndex.setupDimension()
let loadedItems = try! similarityIndex2.loadIndex(name: "TestIndexForLoading")

XCTAssertNotNil(loadedItems)
Expand Down
Loading