added get_num_params and update test cases

.github/workflows/hugging-face.yaml

@@ -12,6 +12,7 @@ jobs:
         with:
           fetch-depth: 0
           lfs: true
+          
       - name: Push to hub
         env:
           HF_TOKEN: ${{ secrets.HF_TOKEN }}

README.md

@@ -65,18 +65,6 @@ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file
 
 Mention any contributors or libraries that you used or were inspired by.
 
----
-
-title: Yume
-emoji: ✨
-colorFrom: blue
-colorTo: green
-sdk: streamlit
-sdk_version: 1.29.0
-app_file: interface.py
-pinned: false
-license: openrail
----
 
 ## Contact
 

sampler.py

@@ -1,7 +0,0 @@
-from yume import Yume
-
-yume = Yume()
-
-yume.load_pretrained()
-
-yume.generate()

sampling.py

@@ -0,0 +1,12 @@
+from .yume import Yume,Config
+
+config = Config()
+
+yume = Yume(config=config)
+
+# Test the quality before loading the pretained
+yume.sample()
+
+yume.load_pretrained()
+
+yume.sample()

tests/test_datasets.py

@@ -1,14 +1,23 @@
 import unittest
-import yume
+from yume.dataset import Trainset 
 
 
 class TestDatasets(unittest.TestCase):
     def test_download(self):
-        pass
+        trainset = Trainset()
+        trainset._load_dataset()
+        assert trainset.texts is not None
+        trainset._tokenize()
+        assert len(trainset.data) > 1
 
     def test_encode(self):
-        pass
-
+        trainset = Trainset()
+        dummy_text = "Hello Human World"
+        trainset.texts = dummy_text
+        trainset._tokenize()
+        assert len(trainset.data) > 1
+        encoded_text = trainset.tokenizer.encode(dummy_text)
+        assert trainset.tokenizer.decode(encoded_text) == dummy_text
 
 if __name__ == "__main__":
     unittest.main()

tests/test_pretrained.py

@@ -1,11 +1,19 @@
 import unittest
+from yume import Yume,Config
 
 
 class TestPretrained(unittest.TestCase):
+    def __init__(self, methodName: str = "runTest") -> None:
+        super().__init__(methodName)
+        self.config = Config()
+        self.yume = Yume(config=self.config)
+    
     def test_download(self):
+        self.yume.load_pretrained()
         pass
 
     def test_generation(self):
+        self.yume.sample()
         pass
 
 

tests/test_tokenizer.py

@@ -1,14 +1,19 @@
 import unittest
-
+from yume import Tokenizer
 
 class TestTokenizer(unittest.TestCase):
+    def __init__(self, methodName: str = "runTest") -> None:
+        super().__init__(methodName)
+        self.tokenizer = Tokenizer()
+        self.dummy_text = "馬鹿なこと言わないでよ"
+    
     def test_encode(self):
         pass
 
     def test_decode(self):
         pass
 
-    def test_equal_result(self):
+    def test_train_encoder(self):
         pass
 
 

training.py

@@ -1,5 +1,21 @@
-from .yume import Yume
+from .yume import Yume,Trainset,Config
 
-yume = Yume()
+config = Config()
 
-yume.train()
+dataset = Trainset()
+
+dataset._load_dataset()
+
+dataset._tokenize(tiktoken=True)
+
+yume = Yume(config)
+
+assert len(dataset.data) > 0
+
+yume.pretrain(dataset.data)
+
+yume.sample()
+
+#optional 
+# yume.huggingface_login("your hf tokens")
+# yume.save_pretrained("yume")

yume/__init__.py

@@ -0,0 +1,4 @@
+from .yume import Yume
+from .dataset import Trainset
+from.tokenizer import Tokenizer
+from .config import Config

yume/config.py

@@ -9,6 +9,7 @@ class Config:
         n_embd=768,
         dropout=0.0,
         bias=True,
+        lr=0.001
     ) -> None:
         self.num_epoch = num_epoch
         self.block_sized = 1024
@@ -18,3 +19,4 @@ class Config:
         self.n_embdd = 768
         self.dropout = 0.0
         self.bias = True
+        self.lr = lr

yume/dataset.py

@@ -1,19 +1,39 @@
 from torch.utils.data import Dataset
 from datasets import load_dataset
 from .tokenizer import Tokenizer
+from .utils import dummy_logger
+
+import tiktoken
 
 
 # TODO setup dataset
 class Trainset(Dataset):
     def __init__(self, batch_size=48):
-        self.loaded_data = load_dataset("zaibutcooler/animanga-vault")
-        self.texts = self.loaded_data["train"]["raw"]
-        self.data = self.loaded_data["train"]["data"]
-        self.tokenizer = Tokenizer()
-        self.tokenizer.load_pretrained()
+        self.texts = None
+        self.data = []
 
     def __len__(self):
         return len(self.data)
-
+    
     def __getitem__(self, index):
+        assert len(self.data) > 10
         return []
+
+
+    def _load_dataset(self,url="zaibutcooler/animanga-vault"):
+        loaded_dataset = load_dataset(url)
+        self.texts = self.loaded_data["train"]["raw"]
+        self.data = self.loaded_data["train"]["data"]
+        dummy_logger("Successfully loaded the dataset")
+    
+    def _tokenize(self,tiktoken=True):
+        if tiktoken:
+            enc = tiktoken.get_encoding("cl100k_base")
+            assert enc.decode(enc.encode("hello world")) == "hello world"
+
+            enc = tiktoken.encoding_for_model("gpt-4")
+            self.tokenizer = enc
+        else:
+            self.tokenizer = Tokenizer()
+            self.tokenizer.load_pretrained()
+        

yume/yume.py

@@ -1,9 +1,10 @@
 import torch
 from torch import nn
 import torch.nn.functional as F
+from huggingface_hub import login
+
 from .config import Config
 from .models import GPT
-from huggingface_hub import login
 from .utils import dummy_logger, training_logger
 
 
@@ -15,9 +16,29 @@ class Yume:
         self.model = GPT(config=config)
         self.config = config
 
-    def train(self):
+    def generate(self):
+        pass
+    
+    def sample(self):
         pass
 
+    def pretrain(self,tokens):
+        lr = self.config.lr
+        num_epochs = self.config.num_epoch
+        
+        
+        pass
+    
+    def fine_tune(self):
+        pass
+
+    def get_num_params(self, non_embedding=True):
+        n_params = sum(p.numel() for p in self.parameters())
+        if non_embedding:
+            n_params -= self.transformer.wpe.weight.numel()
+        dummy_logger(f"parameter count -> {n_params}")
+        return n_params
+    
     def save_pretrained(self, name="yume"):
         self.model.save_pretrained(name)
         self.model.push_to_hub(name)
@@ -30,7 +51,4 @@ class Yume:
     def huggingface_login(self, token):
         assert token is not None
         login(token=token)
-        dummy_logger("Logged in successfully")
-
-    def generate(self):
-        pass
+        dummy_logger("Logged in successfully")