fixed the tokenizer

Author: shivendrra

test.py

@@ -1,44 +1,44 @@
-# import os
-# current_dir = os.path.dirname(os.path.realpath(__file__))
-# os.chdir(current_dir)
-
-# from tokenizer import KMerTokenizer
-
-# tokenizer = KMerTokenizer(k_mers=8)
-
-# with open('training files/file1.txt', 'r', encoding='utf-8') as f:
-#   test_data = f.read().lower()
-#   print("file opened!")
-# f.close()
-# tokenizer.load_model('tokenizer/vocabs/base_4k.json')
-
-# encoded_tokens = tokenizer.encode(test_data)
-# print(encoded_tokens)
-# decoded_tokens = tokenizer.decode(encoded_tokens)
-# print(decoded_tokens)
-# print(f"seq length: {len(test_data)} \ntokens length: {len(decoded_tokens)}")
-# print(test_data == decoded_tokens)
-# print(f"file length: {len(test_data)} \ntokens: {len(encoded_tokens)}")
-# print(f"compression ration: {(len(test_data) / len(encoded_tokens)):.2f}x")
-
 import os
 current_dir = os.path.dirname(os.path.realpath(__file__))
 os.chdir(current_dir)
 
-from tokenizer import PerChar
-tokenizer = PerChar()
+from tokenizer import KMerTokenizer
+
+tokenizer = KMerTokenizer(k_mers=8)
 
 with open('training files/file1.txt', 'r', encoding='utf-8') as f:
-  test_data = f.read()
+  test_data = f.read().lower()
   print("file opened!")
 f.close()
+tokenizer.load_model('tokenizer/vocabs/base_4k.json')
 
 encoded_tokens = tokenizer.encode(test_data)
 print(encoded_tokens)
 decoded_tokens = tokenizer.decode(encoded_tokens)
 print(decoded_tokens)
-
 print(f"seq length: {len(test_data)} \ntokens length: {len(decoded_tokens)}")
 print(test_data == decoded_tokens)
 print(f"file length: {len(test_data)} \ntokens: {len(encoded_tokens)}")
-print(f"compression ration: {(len(test_data) / len(encoded_tokens)):.2f}x")
+print(f"compression ration: {(len(test_data) / len(encoded_tokens)):.2f}x")
+
+# import os
+# current_dir = os.path.dirname(os.path.realpath(__file__))
+# os.chdir(current_dir)
+
+# from tokenizer import PerChar
+# tokenizer = PerChar()
+
+# with open('training files/file1.txt', 'r', encoding='utf-8') as f:
+#   test_data = f.read()
+#   print("file opened!")
+# f.close()
+
+# encoded_tokens = tokenizer.encode(test_data)
+# print(encoded_tokens)
+# decoded_tokens = tokenizer.decode(encoded_tokens)
+# print(decoded_tokens)
+
+# print(f"seq length: {len(test_data)} \ntokens length: {len(decoded_tokens)}")
+# print(test_data == decoded_tokens)
+# print(f"file length: {len(test_data)} \ntokens: {len(encoded_tokens)}")
+# print(f"compression ration: {(len(test_data) / len(encoded_tokens)):.2f}x")

tokenizer/base.py

@@ -1,6 +1,5 @@
 import kmer_c
 import json
-from tqdm import tqdm
 
 class KMerTokenizer:
   def __init__(self, k_mers: int = 4):
@@ -16,7 +15,7 @@ def encode(self, sequence):
 
   def decode(self, encoded_sequence):
     return self.tokenizer.decode(encoded_sequence)
-  
+    
   def save_model(self, model_path):
     vocab_file = f"{model_path}/base_{self.k_mers}k.json"
     with open(vocab_file, 'w') as f:
@@ -30,21 +29,4 @@ def load_model(self, path):
     print("loaded the vocab!")
 
     self.vocab = vocab
-    self.tokenizer.set_vocab(vocab)
-    self.tokenizer.vocab_size = len(vocab)
-
-    self.id_to_token = [None] * self.vocab_size
-    for token, idx in self.vocab.items():
-      self.id_to_token[idx] = token
-
-# if __name__ == "__main__":
-#   tokenizer = KMerTokenizer(k_mers=4)
-#   sequences = ["ATGCGTAC", "GTCAGTAC"]
-#   for sequence in sequences:
-#     print(tokenizer.tokenize_sequence(sequence))
-#     encoded = tokenizer.encode(sequence)
-#     print(encoded)
-#     decoded = tokenizer.decode(encoded)
-#     print(decoded)
-#   tokenizer.save_model("model")
-#   tokenizer.load_model("model/base_4k.json")
+    self.tokenizer.set_vocab(vocab)

tokenizer/csrc/kmer.cpp

@@ -3,7 +3,7 @@
 #include <cmath>
 #include <algorithm>
 
-KMerTokenizer::KMerTokenizer(int k_mers) : k_mers(k_mers) {}
+KMerTokenizer::KMerTokenizer(int k_mers) : k_mers(k_mers), vocab_size(0) {}
 
 std::vector<std::string> KMerTokenizer::tokenize_sequence(const std::string &sequence) {
   std::vector<std::string> kmers;
@@ -20,7 +20,10 @@ std::vector<int> KMerTokenizer::encode(const std::string &sequence) {
     if (token_to_id.find(kmer) != token_to_id.end()) {
       encoded_sequence.push_back(token_to_id[kmer]);
     } else {
-      encoded_sequence.push_back(token_to_id.size() + 1);
+      int new_id = token_to_id.size();
+      token_to_id[kmer] = new_id;
+      id_to_token.push_back(kmer);
+      encoded_sequence.push_back(new_id);
     }
   }
   return encoded_sequence;
@@ -44,6 +47,7 @@ void KMerTokenizer::set_vocab(const std::unordered_map<std::string, int> &vocab)
   for (const auto &pair : vocab) {
     id_to_token[pair.second] = pair.first;
   }
+  vocab_size = vocab.size();
 }
 
 std::unordered_map<std::string, int> KMerTokenizer::get_vocab() {

tokenizer/csrc/kmer.h

@@ -19,6 +19,7 @@ class KMerTokenizer {
   int k_mers;
   std::unordered_map<std::string, int> token_to_id;
   std::vector<std::string> id_to_token;
+  int vocab_size;
 };
 
 #endif