data.py updated to save processed data

src/pipes/data.py

@@ -1,6 +1,7 @@
 import string
 import os
 import random
+import utils
 
 class Dataset:
     def __init__(self, sentences, vocab=None):
@@ -10,7 +11,6 @@ class Dataset:
         vocab = set()
         for sentence in sentences:
             vocab.update(sentence.split())
-        print("Len", len(list(vocab)))
         return sorted(list(vocab))
     
     def remove_punctuation(self, sentence):
@@ -42,22 +42,54 @@ class Dataset:
     
         self.vocab = self.build_vocab(self.sentences)
 
-        if max_length is None:
-            max_length = max(len(sentence.split()) for sentence in self.sentences)
+        self.max_length = max(len(sentence.split()) for sentence in self.sentences)
 
         processed_sentences = []
         for sentence in self.sentences:
             tokens = self.tokenize(sentence)
-            padded_tokens = self.pad_sequence(tokens, max_length)
+            padded_tokens = self.pad_sequence(tokens, self.max_length)
             processed_sentences.append(padded_tokens)
 
-        return processed_sentences
+        data_dict = {}
+        data_dict["max_seq_len"] = self.max_length
+        data_dict["vocab_size"] = len(self.vocab)
+        data_dict["vocab"] = self.vocab
+        return processed_sentences, data_dict
 
+if __name__ == "__main__":
+    # Process
+    # gr
+    sentences = utils.read_txt("E:/bn_multi_tribe_mt/data/raw/gr.txt")
+    dataset = Dataset(sentences)
+    gr, gr_dict = dataset.process()
 
-def load_data(file_path):
-    sentences = []
-    with open(file_path, 'r', encoding='utf-8') as f:
-        sentences = f.readlines()
+    # bn
+    sentences = utils.read_txt("E:/bn_multi_tribe_mt/data/raw/bn.txt")
     dataset = Dataset(sentences)
-    processed_sentences = dataset.process()
-    return processed_sentences
+    bn, bn_dict = dataset.process()
+
+    #Shuffle
+    zipped = list(zip(gr, bn))
+    random.shuffle(zipped)
+    gr, bn = zip(*zipped)
+
+    # Split
+    split_id = int(len(gr) * 0.8)
+
+    gr_train = gr[:split_id]
+    gr_val = gr[split_id:]
+    bn_train = bn[:split_id]
+    bn_val = bn[split_id:]
+
+    # Save
+    gr_dict["train"] = gr_train
+    gr_dict["val"] = gr_val
+    bn_dict["train"] = bn_train
+    bn_dict["val"] = bn_val
+
+    data_dict = {}
+    data_dict["gr"] = gr_dict
+    data_dict["bn"] = bn_dict
+    utils.save_dict("E:/bn_multi_tribe_mt/data/bn_multi_tribe_mt.txt", data_dict)
+
+    # print("Loaded dict: ", utils.load_dict("E:/bn_multi_tribe_mt/data/bn_multi_tribe_mt.txt")["bn"]["val"])

src/pipes/utils.py

@@ -0,0 +1,16 @@
+import json
+
+def save_dict(file_path, my_dict, encoding='utf-8'):
+    with open(file_path, "w", encoding=encoding) as f:
+        json.dump(my_dict, f, ensure_ascii=False)
+
+def load_dict(file_path):
+    with open(file_path, "r", encoding='utf-8') as f:
+        my_dict = json.load(f)
+    return my_dict
+
+def read_txt(file_path):
+    sentences = []
+    with open(file_path, "r", encoding="utf-8") as f:
+        sentences = f.readlines()
+    return sentences