Zai commited on
Commit
2d880e2
·
1 Parent(s): 3fadc1d

updated the training and sampling

Browse files
Files changed (2) hide show
  1. sampling.py +10 -8
  2. training.py +17 -13
sampling.py CHANGED
@@ -1,12 +1,14 @@
1
- from .yume import Yume, Config
 
2
 
3
- config = Config()
 
4
 
5
- yume = Yume(config=config)
 
6
 
7
- # Test the quality before loading the pretained
8
- yume.sample()
9
 
10
- yume.load_pretrained()
11
-
12
- yume.sample()
 
1
+ from yume import Yume
2
+ from yume.config import yume_small
3
 
4
+ # Optional: Create a custom config if needed
5
+ # dummy_config = Config(...)
6
 
7
+ # Initialize the Yume model with a pre-defined small configuration
8
+ yume = Yume(config=yume_small)
9
 
10
+ # Load a pretrained model from the specified path
11
+ yume.load_pretrained('zaibutcooler/yume')
12
 
13
+ # Generate a sample with the prompt '犬とは' (What is a dog?)
14
+ yume.sample('犬とは')
 
training.py CHANGED
@@ -1,21 +1,25 @@
1
- from .yume import Yume, Trainset, Config
 
 
2
 
3
- config = Config()
4
-
5
- dataset = Trainset()
6
 
 
7
  dataset.build_dataset()
8
 
9
- yume = Yume(config)
10
-
11
-
12
 
13
- # assert len(dataset.data) > 0
 
14
 
15
- # yume.pretrain(dataset)
 
16
 
17
- # yume.sample()
 
18
 
19
- # optional
20
- # yume.huggingface_login("your hf tokens")
21
- # yume.save_pretrained("yume")
 
1
+ from yume import Yume
2
+ from yume.dataset import Trainset
3
+ from yume.config import yume_medium, Config
4
 
5
+ # Initialize the dataset with the desired URL
6
+ dataset = Trainset(dataset_url="zaibutcooler/nihon-wiki")
 
7
 
8
+ # Build the dataset
9
  dataset.build_dataset()
10
 
11
+ # Optional: Create a custom config if needed
12
+ # dummy_config = Config(...)
 
13
 
14
+ # Initialize the Yume model with a pre-defined medium configuration
15
+ yume = Yume(config=yume_medium)
16
 
17
+ # Pretrain the model with the dataset
18
+ yume.pretrain(dataset)
19
 
20
+ # Optional: Fine-tune the model with the dataset
21
+ # yume.fine_tune(dataset)
22
 
23
+ # Optional: Upload the model to Hugging Face
24
+ # yume.huggingface_login("your_hf_tokens")
25
+ # yume.save_pretrained("zaibutcooler/yume")