updated the training and sampling

sampling.py

@@ -1,12 +1,14 @@
-from .yume import Yume, Config
+from yume import Yume
+from yume.config import yume_small
 
-config = Config()
+# Optional: Create a custom config if needed
+# dummy_config = Config(...)
 
-yume = Yume(config=config)
+# Initialize the Yume model with a pre-defined small configuration
+yume = Yume(config=yume_small)
 
-# Test the quality before loading the pretained
-yume.sample()
+# Load a pretrained model from the specified path
+yume.load_pretrained('zaibutcooler/yume')
 
-yume.load_pretrained()
-
-yume.sample()
+# Generate a sample with the prompt '犬とは' (What is a dog?)
+yume.sample('犬とは')

training.py

@@ -1,21 +1,25 @@
-from .yume import Yume, Trainset, Config
+from yume import Yume
+from yume.dataset import Trainset
+from yume.config import yume_medium, Config
 
-config = Config()
-
-dataset = Trainset()
+# Initialize the dataset with the desired URL
+dataset = Trainset(dataset_url="zaibutcooler/nihon-wiki")
 
+# Build the dataset
 dataset.build_dataset()
 
-yume = Yume(config)
-
-
+# Optional: Create a custom config if needed
+# dummy_config = Config(...)
 
-# assert len(dataset.data) > 0
+# Initialize the Yume model with a pre-defined medium configuration
+yume = Yume(config=yume_medium)
 
-# yume.pretrain(dataset)
+# Pretrain the model with the dataset
+yume.pretrain(dataset)
 
-# yume.sample()
+# Optional: Fine-tune the model with the dataset
+# yume.fine_tune(dataset)
 
-# optional
-# yume.huggingface_login("your hf tokens")
-# yume.save_pretrained("yume")
+# Optional: Upload the model to Hugging Face
+# yume.huggingface_login("your_hf_tokens")
+# yume.save_pretrained("zaibutcooler/yume")