Adding a constructor for DocTokenizer, and loading `NeuralAttention…

…Lib` and `TextEncodeBase`.
JuliaGenAI · May 31, 2024 · 41f028b · 41f028b
1 parent 72404a8
commit 41f028b
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 0 deletions.
diff --git a/src/ColBERT.jl b/src/ColBERT.jl
@@ -2,6 +2,8 @@ module ColBERT
 using CSV
 using Dates
 using Logging
+using NeuralAttentionlib
+using TextEncodeBase
 using Transformers
 
 # datasets

diff --git a/src/modelling/tokenization/doc_tokenization.jl b/src/modelling/tokenization/doc_tokenization.jl
@@ -1,5 +1,11 @@
 using ...ColBERT: ColBERTConfig
 
 struct DocTokenizer
+    D_marker_token_id::Int
     config::ColBERTConfig
 end
+
+function DocTokenizer(tokenizer::Transformers.TextEncoders.AbstractTransformerTextEncoder, config::ColBERTConfig)
+    D_marker_token_id = TextEncodeBase.lookup(tokenizer.vocab, config.tokenizer_settings.doc_token_id)
+    DocTokenizer(D_marker_token_id, config)
+end