tensorflow tokenizer

tiger.py

@@ -13,22 +13,29 @@ else:
     exit()
 
 
-def process_data(x):
-    x = [item.upper() for item in x]
-    number_of_input = len(x) - GUIDE_LEN + 1
-    input_gens = []
-    for i in range(number_of_input):
-        input_gens.append("".join(x[i:i + GUIDE_LEN]))
-    merged_token = []
-    token_x = [NUCLEOTIDE_TOKENS[item] for item in x]
-    for i in range(number_of_input):
-        merged_token.extend(token_x[i:i + GUIDE_LEN])
-    one_hot_x = tf.one_hot(merged_token, depth=4)
-    model_input_x = tf.reshape(one_hot_x, [-1, GUIDE_LEN * 4])
-    return input_gens, model_input_x
+def process_data(transcript_seq: str):
+
+    # convert to upper case
+    transcript_seq = transcript_seq.upper()
+
+    # get all target sites
+    num_target_sites = len(transcript_seq) - GUIDE_LEN + 1
+    target_seq = [transcript_seq[i:i + GUIDE_LEN] for i in range(num_target_sites)]
+
+    # get one-hot encodings
+    nucleotide_table = tf.lookup.StaticVocabularyTable(
+        initializer=tf.lookup.KeyValueTensorInitializer(
+            keys=tf.constant(list(NUCLEOTIDE_TOKENS.keys()), dtype=tf.string),
+            values=tf.constant(list(NUCLEOTIDE_TOKENS.values()), dtype=tf.int64)),
+        num_oov_buckets=1)
+    target_tokens = nucleotide_table.lookup(tf.stack([list(t) for t in target_seq], axis=0))
+    target_one_hot = tf.reshape(tf.one_hot(target_tokens, depth=4), [num_target_sites, -1])
+
+    return target_seq, target_one_hot
 
 
 def tiger_predict(transcript_seq: str):
+
     # parse transcript sequence into 23-nt target sequences and their one-hot encodings
     target_seq, target_seq_one_hot = process_data(transcript_seq)
 
@@ -42,6 +49,6 @@ def tiger_predict(transcript_seq: str):
 if __name__ == '__main__':
 
     # simple test case
-    transcript_sequence = 'ACGTACGTACGTACGTACGTACGTACGTACGT'
+    transcript_sequence = 'ACGTACGTACGTACGTACGTACGTACGTACGT'.lower()
     df = tiger_predict(transcript_sequence)
     print(df)