Changed initializers of HNSWCorpus to match dictionary

Sources/SwiftNLP/1. Data Collection/HNSWCorpus.swift

@@ -25,19 +25,33 @@ import Foundation
 
 class HNSWCorpus<Scalar: BinaryFloatingPoint & Codable>: SNLPCorpus {
             
-    internal var _documentEncoder: ContextFreeEncoder<Scalar>
+    // internal var _documentEncoder: ContextFreeEncoder<Scalar>
+    internal var _documentEncoder: any SNLPEncoder
     var zeroes: [Scalar]
     var count: Int { 0 }
 
     // typicalNeighbourhoodSize = 20 is a standard benchmark
     var encodedDocuments: DeterministicSampleVectorIndex<[Scalar]>
-
-    init(_documentEncoder: ContextFreeEncoder<Scalar>, typicalNeighborhoodSize: Int = 20) {
-        self._documentEncoder = _documentEncoder
-        zeroes = Array(repeating: Scalar(0), count: 384)
+    
+    init(encoding: ContextFreeEncoder<Scalar>.PreComputedEmbeddings, scalar: Scalar.Type = Double.self,
+         typicalNeighborhoodSize: Int = 20) {
+        _documentEncoder = ContextFreeEncoder(source: encoding)
+        zeroes = _documentEncoder.zeroes as! [Scalar]
         encodedDocuments = DeterministicSampleVectorIndex<[Scalar]>(typicalNeighborhoodSize: typicalNeighborhoodSize)
     }
     
+    init(encoder: any SNLPEncoder, scalar: Scalar.Type = Double.self, typicalNeighborhoodSize: Int = 20) {
+        _documentEncoder = encoder
+        zeroes = _documentEncoder.zeroes as! [Scalar]
+        encodedDocuments = DeterministicSampleVectorIndex<[Scalar]>(typicalNeighborhoodSize: typicalNeighborhoodSize)
+    }
+
+//    init(_documentEncoder: ContextFreeEncoder<Scalar>, typicalNeighborhoodSize: Int = 20) {
+//        self._documentEncoder = _documentEncoder
+//        zeroes = Array(repeating: Scalar(0), count: 384)
+//        encodedDocuments = DeterministicSampleVectorIndex<[Scalar]>(typicalNeighborhoodSize: typicalNeighborhoodSize)
+//    }
+    
     // Decodable conformance
 //    required init(from decoder: Decoder) throws {
 //        let container = try decoder.container(keyedBy: CodingKeys.self)
@@ -48,7 +62,8 @@ class HNSWCorpus<Scalar: BinaryFloatingPoint & Codable>: SNLPCorpus {
     
     @inlinable
     func addUntokenizedDocument(_ document: String) {
-        encodedDocuments.insert((_documentEncoder.encodeSentence(document))) /// as! [Scalar] not needed
+        /// forced unwrap as! [Scalar] is needed when we use SNLPEncoder and not ContextFreeEncoder
+        encodedDocuments.insert((_documentEncoder.encodeSentence(document)) as! [Scalar])
     }
     
 //    @inlinable

Tests/SwiftNLPTests/2. Encoding/HNSWTests.swift

@@ -15,8 +15,9 @@ final class HNSWTests: XCTestCase {
          ]
         
         // let encoder = ContextFreeEncoder<Double>(source: .glove6B50d)
-        let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
-        var corpus = HNSWCorpus(_documentEncoder: encoder)
+        // let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
+        // var corpus = HNSWCorpus(_documentEncoder: encoder)
+        var corpus = HNSWCorpus(encoding: .glove6B50d)
         corpus.addUntokenizedDocuments(docs)
         
         XCTAssert(corpus.encodedDocuments.base.vectors.count == 3)
@@ -54,8 +55,9 @@ final class HNSWTests: XCTestCase {
         ]
         
         // let encoder = ContextFreeEncoder<Double>(source: .glove6B50d)
-        let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
-        var corpus = HNSWCorpus(_documentEncoder: encoder)
+        // let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
+        // var corpus = HNSWCorpus(_documentEncoder: encoder)
+        var corpus = HNSWCorpus(encoding: .glove6B50d)
         corpus.addUntokenizedDocuments(twentyQuotes)
         
         
@@ -79,8 +81,9 @@ final class HNSWTests: XCTestCase {
         let (submissions, _ ): ([Submission],[Data]) = try await loadFromRedditArchive(submissionsData)
         
         // let encoder = ContextFreeEncoder<Double>(source: .glove6B50d)
-        let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
-        var corpus = HNSWCorpus(_documentEncoder: encoder)
+        // let encoder = ContextFreeEncoder<Double>(source: .hnswindex)
+        // var corpus = HNSWCorpus(_documentEncoder: encoder)
+        var corpus = HNSWCorpus(encoding: .glove6B50d)
         
         for submission in submissions {
             if let text = submission.selftext {