Adding the DocTokenizer type, and loading in inside Checkpoint.

JuliaGenAI · May 30, 2024 · 8a6ac59 · 8a6ac59
1 parent b28af5c
commit 8a6ac59
Show file tree

Hide file tree

Showing 5 changed files with 21 additions and 5 deletions.
diff --git a/src/ColBERT.jl b/src/ColBERT.jl
@@ -16,9 +16,10 @@ export RunSettings, TokenizerSettings, ResourceSettings,
        DocSettings, QuerySettings, IndexingSettings,
        SearchSettings, ColBERTConfig
 
-# models
+# models, document/query tokenizers
+include("modelling/tokenization/doc_tokenization.jl")
 include("modelling/checkpoint.jl")
-export BaseColBERT, Checkpoint
+export BaseColBERT, Checkpoint, DocTokenizer
 
 # indexer
 include("indexing.jl")

diff --git a/src/indexing/collection_encoder.jl b/src/indexing/collection_encoder.jl
@@ -1,2 +1,11 @@
+using ..ColBERT: ColBERTConfig
+
 struct CollectionEncoder
+    config::ColBERTConfig
+end
+
+function encode_passages(encoder::CollectionEncoder, passages::Vector{String})
+    @info "Encoding $(length(passages)) passages."
+
+    # TODO: complete this implementation!
 end
diff --git a/src/modelling/checkpoint.jl b/src/modelling/checkpoint.jl
@@ -1,3 +1,5 @@
+using ..ColBERT: DocTokenizer, ColBERTConfig
+
 struct BaseColBERT
     bert::Any
     linear::Any
@@ -18,8 +20,7 @@ function BaseColBERT(checkpoint::String, config::ColBERTConfig)
 end
 
 struct Checkpoint
-
     model::BaseColBERT
-    doc_tokenizer::Any
-    colbert_config::Any
+    doc_tokenizer::DocTokenizer
+    colbert_config::ColBERTConfig
 end
diff --git a/src/modelling/tokenization/doc_tokenization.jl b/src/modelling/tokenization/doc_tokenization.jl
@@ -0,0 +1,5 @@
+using ...ColBERT: ColBERTConfig
+
+struct DocTokenizer
+    config::ColBERTConfig
+end
diff --git a/src/modelling/tokenization/query_tokenization.jl b/src/modelling/tokenization/query_tokenization.jl