word2vec-trainer-pytorch 0.1.0__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -0,0 +1,21 @@
1
+ MIT License
2
+
3
+ Copyright (c) 2026 Abhishek Biswas
4
+
5
+ Permission is hereby granted, free of charge, to any person obtaining a copy
6
+ of this software and associated documentation files (the "Software"), to deal
7
+ in the Software without restriction, including without limitation the rights
8
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
9
+ copies of the Software, and to permit persons to whom the Software is
10
+ furnished to do so, subject to the following conditions:
11
+
12
+ The above copyright notice and this permission notice shall be included in all
13
+ copies or substantial portions of the Software.
14
+
15
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
16
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
17
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
18
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
19
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
20
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
21
+ SOFTWARE.
@@ -0,0 +1,166 @@
1
+ Metadata-Version: 2.4
2
+ Name: word2vec-trainer-pytorch
3
+ Version: 0.1.0
4
+ Summary: A PyTorch implementation of Word2Vec for learning and experimentation
5
+ Author-email: Abhishek Biswas <rohan.abiswas@gmail.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2026 Abhishek Biswas
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Repository, https://github.com/AbhishekBiswas12/word2vec_trainer
29
+ Project-URL: Issues, https://github.com/AbhishekBiswas12/word2vec_trainer/issues
30
+ Requires-Python: >=3.10
31
+ Description-Content-Type: text/markdown
32
+ License-File: LICENSE
33
+ Requires-Dist: torch>=2.0
34
+ Requires-Dist: numpy>=1.24
35
+ Requires-Dist: pandas>=2.0
36
+ Requires-Dist: tqdm>=4.64
37
+ Requires-Dist: matplotlib>=3.8.0
38
+ Requires-Dist: pytest>=9.1.0
39
+ Dynamic: license-file
40
+
41
+ # Word2Vec Trainer
42
+
43
+ A PyTorch implementation of **Word2Vec** built as a learning project to understand how word embeddings are trained.
44
+
45
+ ## About
46
+
47
+ I wanted to understand Word2Vec beyond the API, so I implemented the training pipeline myself and packaged it as a Python library. I built this project while learning about **Word2Vec, Skip-gram, negative sampling, and word embeddings**.
48
+
49
+ The goal was to implement the main components myself and understand how they work rather than treating Word2Vec as a black box.
50
+
51
+ The project includes:
52
+
53
+ - Text preprocessing
54
+ - Vocabulary creation
55
+ - Training pair generation
56
+ - Skip-gram model
57
+ - Negative sampling
58
+ - Model training with PyTorch
59
+
60
+ ## Project Structure
61
+
62
+ ```text
63
+ word2vec_trainer/
64
+
65
+ ├── pyproject.toml
66
+ ├── README.md
67
+ ├── LICENSE
68
+
69
+ ├── src/
70
+ │ └── word2vec_trainer/
71
+ │ ├── dataset_loader.py
72
+ │ ├── word2vec_model.py
73
+ │ ├── preprocessing.py
74
+ │ └── trainer.py
75
+
76
+ ├── tests/
77
+ │ ├── test_dataset.py
78
+ │ ├── test_preprocessing.py
79
+ │ ├── test_model.py
80
+ │ └── test_trainer.py
81
+
82
+ └── examples/
83
+ └── basic_training.py
84
+ ```
85
+
86
+ ## Training Pipeline
87
+
88
+ ```text
89
+ Text without commas or special characters
90
+
91
+ Preprocessing
92
+
93
+ Vocabulary
94
+
95
+ Training Pairs
96
+
97
+ Negative Sampling
98
+
99
+ Skip-gram Model
100
+
101
+ Learned Word Embeddings
102
+ ```
103
+
104
+ ## Installation
105
+
106
+ Clone the repository:
107
+
108
+ ```bash
109
+ git clone https://github.com/AbhishekBiswas12/word2vec_trainer.git
110
+ cd word2vec_trainer
111
+ ```
112
+
113
+ Install in editable mode:
114
+
115
+ ```bash
116
+ python -m pip install -e .
117
+ ```
118
+
119
+ ## Usage
120
+
121
+ A basic training example is available in:
122
+
123
+ ```text
124
+ examples/basic_training.py
125
+ ```
126
+
127
+ The public API is still evolving as the project develops.
128
+
129
+ ## Testing
130
+
131
+ Run the test suite with:
132
+
133
+ ```bash
134
+ pytest
135
+ ```
136
+
137
+ ## AI-Assisted Development
138
+
139
+ This project was built primarily as a learning exercise.
140
+
141
+ During development, I occasionally used **ChatGPT** and **Google Colab's coding assistant** when I got stuck with implementation details, debugging, or understanding concepts.
142
+
143
+ I used these tools as learning and development assistance while working to understand the underlying implementation.
144
+
145
+ ## Status
146
+
147
+ 🚧 **Work in progress**
148
+
149
+ Future improvements include:
150
+
151
+ - Improve the public API
152
+ - Expand test coverage
153
+ - Add more examples
154
+ - Add embedding evaluation and visualization
155
+ - Benchmark different training configurations
156
+ - Publish the package to PyPI
157
+
158
+ ## References
159
+
160
+ - [Efficient Estimation of Word Representations in Vector Space](https://arxiv.org/abs/1301.3781)
161
+ - [Distributed Representations of Words and Phrases and their Compositionality](https://arxiv.org/abs/1310.4546)
162
+ - [The Illustrated Word2Vec](https://jalammar.github.io/illustrated-word2vec/)
163
+
164
+ ## License
165
+
166
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,126 @@
1
+ # Word2Vec Trainer
2
+
3
+ A PyTorch implementation of **Word2Vec** built as a learning project to understand how word embeddings are trained.
4
+
5
+ ## About
6
+
7
+ I wanted to understand Word2Vec beyond the API, so I implemented the training pipeline myself and packaged it as a Python library. I built this project while learning about **Word2Vec, Skip-gram, negative sampling, and word embeddings**.
8
+
9
+ The goal was to implement the main components myself and understand how they work rather than treating Word2Vec as a black box.
10
+
11
+ The project includes:
12
+
13
+ - Text preprocessing
14
+ - Vocabulary creation
15
+ - Training pair generation
16
+ - Skip-gram model
17
+ - Negative sampling
18
+ - Model training with PyTorch
19
+
20
+ ## Project Structure
21
+
22
+ ```text
23
+ word2vec_trainer/
24
+
25
+ ├── pyproject.toml
26
+ ├── README.md
27
+ ├── LICENSE
28
+
29
+ ├── src/
30
+ │ └── word2vec_trainer/
31
+ │ ├── dataset_loader.py
32
+ │ ├── word2vec_model.py
33
+ │ ├── preprocessing.py
34
+ │ └── trainer.py
35
+
36
+ ├── tests/
37
+ │ ├── test_dataset.py
38
+ │ ├── test_preprocessing.py
39
+ │ ├── test_model.py
40
+ │ └── test_trainer.py
41
+
42
+ └── examples/
43
+ └── basic_training.py
44
+ ```
45
+
46
+ ## Training Pipeline
47
+
48
+ ```text
49
+ Text without commas or special characters
50
+
51
+ Preprocessing
52
+
53
+ Vocabulary
54
+
55
+ Training Pairs
56
+
57
+ Negative Sampling
58
+
59
+ Skip-gram Model
60
+
61
+ Learned Word Embeddings
62
+ ```
63
+
64
+ ## Installation
65
+
66
+ Clone the repository:
67
+
68
+ ```bash
69
+ git clone https://github.com/AbhishekBiswas12/word2vec_trainer.git
70
+ cd word2vec_trainer
71
+ ```
72
+
73
+ Install in editable mode:
74
+
75
+ ```bash
76
+ python -m pip install -e .
77
+ ```
78
+
79
+ ## Usage
80
+
81
+ A basic training example is available in:
82
+
83
+ ```text
84
+ examples/basic_training.py
85
+ ```
86
+
87
+ The public API is still evolving as the project develops.
88
+
89
+ ## Testing
90
+
91
+ Run the test suite with:
92
+
93
+ ```bash
94
+ pytest
95
+ ```
96
+
97
+ ## AI-Assisted Development
98
+
99
+ This project was built primarily as a learning exercise.
100
+
101
+ During development, I occasionally used **ChatGPT** and **Google Colab's coding assistant** when I got stuck with implementation details, debugging, or understanding concepts.
102
+
103
+ I used these tools as learning and development assistance while working to understand the underlying implementation.
104
+
105
+ ## Status
106
+
107
+ 🚧 **Work in progress**
108
+
109
+ Future improvements include:
110
+
111
+ - Improve the public API
112
+ - Expand test coverage
113
+ - Add more examples
114
+ - Add embedding evaluation and visualization
115
+ - Benchmark different training configurations
116
+ - Publish the package to PyPI
117
+
118
+ ## References
119
+
120
+ - [Efficient Estimation of Word Representations in Vector Space](https://arxiv.org/abs/1301.3781)
121
+ - [Distributed Representations of Words and Phrases and their Compositionality](https://arxiv.org/abs/1310.4546)
122
+ - [The Illustrated Word2Vec](https://jalammar.github.io/illustrated-word2vec/)
123
+
124
+ ## License
125
+
126
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,31 @@
1
+ [build-system]
2
+ requires = ["setuptools>=68"]
3
+ build-backend = "setuptools.build_meta"
4
+
5
+ [project]
6
+ name = "word2vec-trainer-pytorch"
7
+ version = "0.1.0"
8
+ description = "A PyTorch implementation of Word2Vec for learning and experimentation"
9
+ readme = "README.md"
10
+ requires-python = ">=3.10"
11
+ license = { file = "LICENSE" }
12
+
13
+ authors = [
14
+ { name = "Abhishek Biswas", email = "rohan.abiswas@gmail.com" }
15
+ ]
16
+
17
+ dependencies = [
18
+ "torch>=2.0",
19
+ "numpy>=1.24",
20
+ "pandas>=2.0",
21
+ "tqdm>=4.64",
22
+ "matplotlib>=3.8.0",
23
+ "pytest>=9.1.0"
24
+ ]
25
+
26
+ [project.urls]
27
+ Repository = "https://github.com/AbhishekBiswas12/word2vec_trainer"
28
+ Issues = "https://github.com/AbhishekBiswas12/word2vec_trainer/issues"
29
+
30
+ [tool.setuptools.packages.find]
31
+ where = ["src"]
@@ -0,0 +1,4 @@
1
+ [egg_info]
2
+ tag_build =
3
+ tag_date = 0
4
+
@@ -0,0 +1 @@
1
+ __version__ = "0.1.0"
@@ -0,0 +1,80 @@
1
+ import torch
2
+ from torch.utils.data import IterableDataset
3
+ import numpy as np
4
+ import math
5
+ from tqdm.auto import tqdm
6
+
7
+ class DatasetLoader(IterableDataset):
8
+ total_batches = 0
9
+ def __init__(self, file_path, batch_size, path_to_neg_dist="neg_distribution.npy", neg_num=5):
10
+ self.file_path = file_path
11
+ self.batch_size = batch_size
12
+ self.neg_num = neg_num
13
+ with open(file_path) as f:
14
+ self.length = sum(1 for _ in f) - 1
15
+ self.neg_dist = np.load(path_to_neg_dist)
16
+ self.total_batches=math.ceil((self.length * (1 + self.neg_num))/self.batch_size)
17
+
18
+ def __len__(self):
19
+ return self.length * (1+self.neg_num)
20
+
21
+ def gen_negatives(self, contexts, targets, batch):
22
+ negatives = np.random.choice(
23
+ len(self.neg_dist),
24
+ size=(len(contexts), self.neg_num),
25
+ p=self.neg_dist
26
+ )
27
+ for n, t, c in zip(negatives, targets, contexts):
28
+ for i in range(len(n)):
29
+ while n[i] == t:
30
+ n[i] = np.random.choice(
31
+ len(self.neg_dist),
32
+ p=self.neg_dist
33
+ )
34
+ batch.extend(
35
+ (c, neg, 0) for neg in n
36
+ )
37
+ return batch
38
+
39
+ def __iter__(self):
40
+ with open(self.file_path) as f:
41
+ next(f) # skip header
42
+ np.random.seed(42)
43
+ batch = []
44
+ targets = []
45
+ contexts = []
46
+
47
+ pbar = tqdm(
48
+ total=math.ceil(
49
+ (self.length * (1 + self.neg_num))/self.batch_size
50
+ ),
51
+ desc="Loading batches",
52
+ unit="batches"
53
+ )
54
+
55
+ for line in f:
56
+ context, target, label = line.strip().split(',')
57
+ context = int(context)
58
+ target = int(target)
59
+ label = int(label)
60
+ contexts.append(context)
61
+ targets.append(target)
62
+ batch.append((context, target, label))
63
+ if len(contexts)==10000:
64
+ batch = self.gen_negatives(contexts, targets, batch)
65
+ contexts = []
66
+ targets = []
67
+ while len(batch) >= self.batch_size:
68
+ pbar.update(1)
69
+ yield torch.tensor(
70
+ batch[:self.batch_size]
71
+ )
72
+
73
+ batch = batch[self.batch_size:]
74
+ if contexts:
75
+ batch = self.gen_negatives(contexts, targets, batch)
76
+
77
+ if batch:
78
+ pbar.update(1)
79
+ yield torch.tensor(batch)
80
+ pbar.close()
@@ -0,0 +1,161 @@
1
+ import pandas as pd
2
+ from tqdm.auto import tqdm
3
+ import math
4
+ import random
5
+ import numpy as np
6
+ import csv
7
+
8
+ class Preprocessing:
9
+ vocabulary_size = None
10
+ def __init__(self,
11
+ cleaned_data_file_path = 'cleaned_data.txt',
12
+ subsampled_file_path = 'subsampled_data.txt',
13
+ vocabulary_file = 'vocabulary.csv',
14
+ pos_example = 'pos_examples.csv',
15
+ train_path = 'train_w2v.csv',
16
+ train_ratio = 0.95,
17
+ valid_path = 'valid_w2v.csv',
18
+ valid_ratio = 0.025,
19
+ test_path = 'test_w2v.csv',
20
+ test_ratio = 0.025,
21
+ window = 5):
22
+ self.cleaned_data_file_path = cleaned_data_file_path
23
+ self.subsampled_file_path = subsampled_file_path
24
+ self.vocabulary_file = vocabulary_file
25
+ self.pos_examples = pos_example
26
+ self.window = window
27
+ self.train_path = train_path
28
+ self.train_ratio = train_ratio
29
+ self.valid_path = valid_path
30
+ self.valid_ratio = valid_ratio
31
+ self.test_path = test_path
32
+ self.test_ratio = test_ratio
33
+ with open(self.pos_examples, "w", newline="", encoding="utf-8") as f:
34
+ writer = csv.writer(f)
35
+ # Write header
36
+ writer.writerow(('int', 'target', 'pred'))
37
+
38
+ def run(self):
39
+ count_word = self.vocabulary_subsampling()
40
+ self.negative_distribution_calc()
41
+ self.positive_example_gen(count_word)
42
+ print('Generated Positive examples.')
43
+ print('Splitting data into train, test and valid sets.')
44
+ self.train_test_val_split()
45
+
46
+ def vocabulary_subsampling(self):
47
+ cleaned_text = open(self.cleaned_data_file_path, 'r').read()
48
+ print('Creating Vocabulary...')
49
+ vocabulary = pd.DataFrame(pd.DataFrame(cleaned_text.split(" ")).value_counts())[pd.DataFrame(pd.DataFrame(cleaned_text.split(" ")).value_counts())['count'] >= 2]
50
+ vocabulary.reset_index(inplace=True)
51
+ print('Columns of vocabulary:', vocabulary.columns)
52
+ vocabulary.loc[len(vocabulary)] = [r'<UNK>', 0]
53
+ vocabulary.rename(columns={0: 'Word'}, inplace=True)
54
+ print("Created Vocabulary.")
55
+ vocabulary.to_csv(self.vocabulary_file)
56
+ self.vocabulary_size = len(vocabulary.index)
57
+ print('Saved vocabulary.')
58
+
59
+ # subsampling
60
+ word_int = {v:k for k, v in vocabulary['Word'].to_dict().items()}
61
+ word_freqs = vocabulary["count"].to_dict()
62
+ cleaned_text_int = []
63
+ c=0
64
+ unk = 0
65
+ for word in tqdm(cleaned_text.split(" "), desc='Converting words to integers'):
66
+ if word in word_int.keys():
67
+ cleaned_text_int.append(word_int[word])
68
+ else:
69
+ cleaned_text_int.append(word_int[r'<UNK>'])
70
+ unk+=1
71
+ c+=1
72
+
73
+ t = 1e-5
74
+ keep_probs = {}
75
+ for word, freq in word_freqs.items():
76
+ if word == len(vocabulary)-1:
77
+ keep_prob = (math.sqrt((unk/c) / t) + 1) * (t / (unk/c))
78
+ else:
79
+ keep_prob = (math.sqrt((freq/c) / t) + 1) * (t / (freq/c))
80
+ keep_prob = min(1.0, keep_prob)
81
+ keep_probs[word] = keep_prob
82
+
83
+ count_word = 0
84
+ with open(self.subsampled_file_path, 'w') as f:
85
+ for i in tqdm(cleaned_text_int, desc="Subsampling common words"):
86
+ if random.random() >= keep_probs[i]:
87
+ continue
88
+ f.write(str(i) + ' ')
89
+ count_word+=1
90
+ print('Subsampling finished')
91
+ return count_word
92
+
93
+ def negative_distribution_calc(self):
94
+ text = open(self.subsampled_file_path, 'r').read().split(" ")[:-1]
95
+ text = [int(num) for num in text]
96
+ counts = np.bincount(text)
97
+ scaled_counts = counts**0.75
98
+ neg_dist = (scaled_counts/np.sum(scaled_counts))
99
+ np.save('neg_distribution_test.npy', neg_dist)
100
+
101
+ def positive_example_gen(self, counts):
102
+ sampled_text = open(self.subsampled_file_path, 'r').read().split(" ")[:-1]
103
+ sampled_text = [int(x) for x in sampled_text]
104
+ window = self.window
105
+ pairs = {
106
+ 'int': [],
107
+ 'target': [],
108
+ 'pred': []
109
+ }
110
+ for i in tqdm(range(self.window, counts-window-1), desc='Computing positive pairs'):
111
+ dic = {
112
+ 'int': [sampled_text[i]]*(window*2),
113
+ 'target':sampled_text[i-window:i] + sampled_text[i+1:i+window+1],
114
+ 'pred': [1]*(window*2)
115
+ }
116
+ pairs['int'].extend(dic['int'])
117
+ pairs['target'].extend(dic['target'])
118
+ pairs['pred'].extend(dic['pred'])
119
+
120
+ with open(self.pos_examples, "a", newline="", encoding="utf-8") as f:
121
+ writer = csv.writer(f)
122
+ writer.writerows(zip(*pairs.values()))
123
+ pairs = {
124
+ 'int': [],
125
+ 'target': [],
126
+ 'pred': []
127
+ }
128
+
129
+ def train_test_val_split(self):
130
+ train_count = 0
131
+ test_count = 0
132
+ valid_count = 0
133
+
134
+ with open(self.pos_examples, "r", newline="", encoding="utf-8") as infile, \
135
+ open(self.train_path, "w", newline="", encoding="utf-8") as train_out, \
136
+ open(self.valid_path, "w", newline="", encoding="utf-8") as valid_out, \
137
+ open(self.test_path, "w", newline="", encoding="utf-8") as test_out:
138
+ reader = csv.reader(infile)
139
+ train_writer = csv.writer(train_out)
140
+ valid_writer = csv.writer(valid_out)
141
+ test_writer = csv.writer(test_out)
142
+
143
+ header = next(reader)
144
+
145
+ train_writer.writerow(header)
146
+ valid_writer.writerow(header)
147
+ test_writer.writerow(header)
148
+ for row in tqdm(reader, desc='Splitting data into train, test and valid sets'):
149
+ r = random.random()
150
+ if r < self.train_ratio:
151
+ train_writer.writerow(row)
152
+ train_count+=1
153
+ elif r < self.train_ratio + self.valid_ratio:
154
+ valid_writer.writerow(row)
155
+ valid_count+=1
156
+ else:
157
+ test_writer.writerow(row)
158
+ test_count+=1
159
+ print('Final Train count: ',train_count)
160
+ print('Final Valid count: ', valid_count)
161
+ print('Final Test count: ', test_count)
@@ -0,0 +1,162 @@
1
+ import torch
2
+ import torch.nn.functional as F
3
+ import torch.optim as op
4
+ import numpy as np
5
+ import math
6
+ from tqdm.auto import tqdm
7
+ from word2vec_trainer.dataset_loader import DatasetLoader
8
+ from word2vec_trainer.word2vec_model import Word2Vec
9
+
10
+ class Trainer:
11
+ def __init__(
12
+ self,
13
+ vocab_size = 508205,
14
+ embedding_dim = 300,
15
+ batch_size_train = 1500000,
16
+ batch_size_val = 50000,
17
+ new_model = True,
18
+ path_to_train = "train.csv",
19
+ path_to_val = "val_csv",
20
+ path_to_saved_model = "",
21
+ path_to_neg_dist = "neg_distribution.npy",
22
+ epochs = 50,
23
+ lr = 0.01
24
+ ):
25
+ # assigning variable values
26
+ self.vocab_size = vocab_size
27
+ self.embedding_dim = embedding_dim
28
+ self.batch_size_train = batch_size_train
29
+ self.batch_size_val = batch_size_val
30
+ self.new_model = new_model
31
+ self.path_to_train = path_to_train
32
+ self.path_to_val = path_to_val
33
+ self.path_to_saved_model = path_to_saved_model if path_to_saved_model != "" else 'word2vec_model.bin'
34
+ self.path_to_neg_dist = path_to_neg_dist
35
+ self.epochs = epochs
36
+
37
+ # creating dataset and model objects
38
+ self.train = DatasetLoader(
39
+ path_to_train,
40
+ batch_size_train,
41
+ path_to_neg_dist=self.path_to_neg_dist,
42
+ neg_num=2)
43
+ self.val = DatasetLoader(
44
+ path_to_val,
45
+ batch_size_val,
46
+ path_to_neg_dist=self.path_to_neg_dist,
47
+ neg_num=2)
48
+ self.model = self.model_creation()
49
+ self.model.requires_grad_(True)
50
+
51
+ # Learning rate, optimizer, loss selection -> TODO: assign from user choice
52
+ self.lr = lr
53
+ self.optimizer = op.Adagrad(self.model.parameters(), lr=self.lr)
54
+ self.loss = torch.nn.BCEWithLogitsLoss()
55
+
56
+ def run(self):
57
+ self.train_model()
58
+
59
+
60
+ def model_creation(self):
61
+ if self.new_model:
62
+ print("Creating train_losses.txt")
63
+ open('train_losses.txt', 'w').close()
64
+ print("Creating val_losses.txt")
65
+ open('val_losses.txt', 'w').close()
66
+ return Word2Vec(self.vocab_size, self.embedding_dim)
67
+ else:
68
+ return torch.load(self.path_to_saved_model, weights_only=False)
69
+
70
+ def validate_model(self, device='cpu'):
71
+ val_loss = 0.0
72
+ val_size = 0
73
+
74
+ train_pos_probs = 0.0
75
+ train_neg_probs = 0.0
76
+ train_pos_size = 0
77
+ train_neg_size = 0
78
+ val_pos_probs = 0.0
79
+ val_neg_probs = 0.0
80
+ val_pos_size = 0
81
+ val_neg_size = 0
82
+
83
+ train_loss = 0.0
84
+ train_size = 0
85
+ self.model.eval()
86
+
87
+ with torch.no_grad():
88
+ for x in tqdm(self.train, desc=f"Computing Train Loss", total=self.train.total_batches):
89
+ context = x[:, 0].to(device)
90
+ target = x[:, 1].to(device)
91
+ labels = x[:, 2].to(device)
92
+ scores = self.model(context, target)
93
+ probs = torch.sigmoid(scores)
94
+
95
+ train_pos_probs += probs[labels == 1].sum().item()
96
+ train_neg_probs += probs[labels == 0].sum().item()
97
+
98
+ train_pos_size += probs[labels==1].size()[0]
99
+ train_neg_size += probs[labels==0].size()[0]
100
+
101
+ l = self.loss(scores, labels.float())
102
+ train_loss += l.item() * context.shape[0]
103
+ train_size += context.shape[0]
104
+
105
+ for y in tqdm(self.val, desc=f"Computing Validation Loss", total=self.val.total_batches):
106
+ context = y[:, 0].to(device)
107
+ target = y[:, 1].to(device)
108
+ labels = y[:, 2].to(device)
109
+
110
+ scores = self.model(context, target)
111
+ probs = torch.sigmoid(scores)
112
+
113
+ val_pos_probs += probs[labels == 1].sum().item()
114
+ val_neg_probs += probs[labels == 0].sum().item()
115
+
116
+ val_pos_size += probs[labels==1].size()[0]
117
+ val_neg_size += probs[labels==0].size()[0]
118
+
119
+ l = self.loss(scores, labels.float())
120
+ val_loss += l.item() * context.shape[0]
121
+ val_size += context.shape[0]
122
+
123
+ train_loss /= train_size
124
+ val_loss /= val_size
125
+
126
+ train_pos_probs /= train_pos_size
127
+ train_neg_probs /= train_neg_size
128
+
129
+ val_pos_probs /= val_pos_size
130
+ val_neg_probs /= val_neg_size
131
+
132
+ self.model.train()
133
+ return train_loss, train_pos_probs, train_neg_probs, val_loss, val_pos_probs, val_neg_probs
134
+
135
+ def train_model(self):
136
+ epoch = len(open('train_losses.txt', 'r').readlines())
137
+ print("starting training...")
138
+ for epoch in range(epoch, self.epochs):
139
+ self.optimizer.zero_grad()
140
+ for x in tqdm(self.train, desc=f"Epoch {epoch+1}", total=self.train.total_batches):
141
+ context = x[:, 0]
142
+ target = x[:, 1]
143
+ labels = x[:, 2]
144
+
145
+ batch_size = target.shape[0]
146
+
147
+ # Forward
148
+ scores = self.model(context, target)
149
+
150
+ l = self.loss(scores, labels.float())
151
+
152
+ # Backprop
153
+ l.backward()
154
+ self.optimizer.step()
155
+ train_loss, train_pos_probs, train_neg_probs, val_loss, val_pos_probs, val_neg_probs = self.validate_model()
156
+ with open('train_losses.txt', 'a') as losses, open('val_losses.txt', 'a') as val_losses:
157
+ losses.write(f"{train_loss}\t{train_pos_probs}\t{train_neg_probs}\n")
158
+ val_losses.write(f"{val_loss}\t{val_pos_probs}\t{val_neg_probs}\n")
159
+ print(f"Epoch {epoch+1}, Loss: {train_loss}, Positive probabs: {train_pos_probs}, Negative probabs: {train_neg_probs}\n Val_loss: {val_loss}, Positive probabs {val_pos_probs}, Negative probabs: {val_neg_probs}")
160
+ if self.path_to_saved_model != '':
161
+ torch.save(self.model, self.path_to_saved_model)
162
+ print(f"Model saved after {epoch+1} epochs")
@@ -0,0 +1,13 @@
1
+ import torch
2
+
3
+ class Word2Vec(torch.nn.Module):
4
+ def __init__(self, vocab_size, embedding_dim):
5
+ super().__init__()
6
+ self.input_embedding_layer = torch.nn.Embedding(vocab_size, embedding_dim)
7
+ self.output_embedding_layer = torch.nn.Embedding(vocab_size, embedding_dim)
8
+
9
+ def forward(self, x, y):
10
+ e1 = self.input_embedding_layer(x)
11
+ e2 = self.output_embedding_layer(y)
12
+
13
+ return (e1 * e2).sum(dim=1)
@@ -0,0 +1,166 @@
1
+ Metadata-Version: 2.4
2
+ Name: word2vec-trainer-pytorch
3
+ Version: 0.1.0
4
+ Summary: A PyTorch implementation of Word2Vec for learning and experimentation
5
+ Author-email: Abhishek Biswas <rohan.abiswas@gmail.com>
6
+ License: MIT License
7
+
8
+ Copyright (c) 2026 Abhishek Biswas
9
+
10
+ Permission is hereby granted, free of charge, to any person obtaining a copy
11
+ of this software and associated documentation files (the "Software"), to deal
12
+ in the Software without restriction, including without limitation the rights
13
+ to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14
+ copies of the Software, and to permit persons to whom the Software is
15
+ furnished to do so, subject to the following conditions:
16
+
17
+ The above copyright notice and this permission notice shall be included in all
18
+ copies or substantial portions of the Software.
19
+
20
+ THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
21
+ IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
22
+ FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
23
+ AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
24
+ LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
25
+ OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
26
+ SOFTWARE.
27
+
28
+ Project-URL: Repository, https://github.com/AbhishekBiswas12/word2vec_trainer
29
+ Project-URL: Issues, https://github.com/AbhishekBiswas12/word2vec_trainer/issues
30
+ Requires-Python: >=3.10
31
+ Description-Content-Type: text/markdown
32
+ License-File: LICENSE
33
+ Requires-Dist: torch>=2.0
34
+ Requires-Dist: numpy>=1.24
35
+ Requires-Dist: pandas>=2.0
36
+ Requires-Dist: tqdm>=4.64
37
+ Requires-Dist: matplotlib>=3.8.0
38
+ Requires-Dist: pytest>=9.1.0
39
+ Dynamic: license-file
40
+
41
+ # Word2Vec Trainer
42
+
43
+ A PyTorch implementation of **Word2Vec** built as a learning project to understand how word embeddings are trained.
44
+
45
+ ## About
46
+
47
+ I wanted to understand Word2Vec beyond the API, so I implemented the training pipeline myself and packaged it as a Python library. I built this project while learning about **Word2Vec, Skip-gram, negative sampling, and word embeddings**.
48
+
49
+ The goal was to implement the main components myself and understand how they work rather than treating Word2Vec as a black box.
50
+
51
+ The project includes:
52
+
53
+ - Text preprocessing
54
+ - Vocabulary creation
55
+ - Training pair generation
56
+ - Skip-gram model
57
+ - Negative sampling
58
+ - Model training with PyTorch
59
+
60
+ ## Project Structure
61
+
62
+ ```text
63
+ word2vec_trainer/
64
+
65
+ ├── pyproject.toml
66
+ ├── README.md
67
+ ├── LICENSE
68
+
69
+ ├── src/
70
+ │ └── word2vec_trainer/
71
+ │ ├── dataset_loader.py
72
+ │ ├── word2vec_model.py
73
+ │ ├── preprocessing.py
74
+ │ └── trainer.py
75
+
76
+ ├── tests/
77
+ │ ├── test_dataset.py
78
+ │ ├── test_preprocessing.py
79
+ │ ├── test_model.py
80
+ │ └── test_trainer.py
81
+
82
+ └── examples/
83
+ └── basic_training.py
84
+ ```
85
+
86
+ ## Training Pipeline
87
+
88
+ ```text
89
+ Text without commas or special characters
90
+
91
+ Preprocessing
92
+
93
+ Vocabulary
94
+
95
+ Training Pairs
96
+
97
+ Negative Sampling
98
+
99
+ Skip-gram Model
100
+
101
+ Learned Word Embeddings
102
+ ```
103
+
104
+ ## Installation
105
+
106
+ Clone the repository:
107
+
108
+ ```bash
109
+ git clone https://github.com/AbhishekBiswas12/word2vec_trainer.git
110
+ cd word2vec_trainer
111
+ ```
112
+
113
+ Install in editable mode:
114
+
115
+ ```bash
116
+ python -m pip install -e .
117
+ ```
118
+
119
+ ## Usage
120
+
121
+ A basic training example is available in:
122
+
123
+ ```text
124
+ examples/basic_training.py
125
+ ```
126
+
127
+ The public API is still evolving as the project develops.
128
+
129
+ ## Testing
130
+
131
+ Run the test suite with:
132
+
133
+ ```bash
134
+ pytest
135
+ ```
136
+
137
+ ## AI-Assisted Development
138
+
139
+ This project was built primarily as a learning exercise.
140
+
141
+ During development, I occasionally used **ChatGPT** and **Google Colab's coding assistant** when I got stuck with implementation details, debugging, or understanding concepts.
142
+
143
+ I used these tools as learning and development assistance while working to understand the underlying implementation.
144
+
145
+ ## Status
146
+
147
+ 🚧 **Work in progress**
148
+
149
+ Future improvements include:
150
+
151
+ - Improve the public API
152
+ - Expand test coverage
153
+ - Add more examples
154
+ - Add embedding evaluation and visualization
155
+ - Benchmark different training configurations
156
+ - Publish the package to PyPI
157
+
158
+ ## References
159
+
160
+ - [Efficient Estimation of Word Representations in Vector Space](https://arxiv.org/abs/1301.3781)
161
+ - [Distributed Representations of Words and Phrases and their Compositionality](https://arxiv.org/abs/1310.4546)
162
+ - [The Illustrated Word2Vec](https://jalammar.github.io/illustrated-word2vec/)
163
+
164
+ ## License
165
+
166
+ MIT License. See [LICENSE](LICENSE) for details.
@@ -0,0 +1,17 @@
1
+ LICENSE
2
+ README.md
3
+ pyproject.toml
4
+ src/word2vec_trainer/__init__.py
5
+ src/word2vec_trainer/dataset_loader.py
6
+ src/word2vec_trainer/preprocessing.py
7
+ src/word2vec_trainer/trainer.py
8
+ src/word2vec_trainer/word2vec_model.py
9
+ src/word2vec_trainer_pytorch.egg-info/PKG-INFO
10
+ src/word2vec_trainer_pytorch.egg-info/SOURCES.txt
11
+ src/word2vec_trainer_pytorch.egg-info/dependency_links.txt
12
+ src/word2vec_trainer_pytorch.egg-info/requires.txt
13
+ src/word2vec_trainer_pytorch.egg-info/top_level.txt
14
+ tests/test_dataset.py
15
+ tests/test_model.py
16
+ tests/test_preprocessing.py
17
+ tests/test_trainer.py
@@ -0,0 +1,6 @@
1
+ torch>=2.0
2
+ numpy>=1.24
3
+ pandas>=2.0
4
+ tqdm>=4.64
5
+ matplotlib>=3.8.0
6
+ pytest>=9.1.0
@@ -0,0 +1,10 @@
1
+ from word2vec_trainer.dataset_loader import DatasetLoader
2
+
3
+ def test_dataset_loader():
4
+ dataset = DatasetLoader('tests/test_inputs/train_w2v.csv', 10, 'tests/test_inputs/neg_distribution_test.npy')
5
+
6
+ assert dataset.file_path == 'tests/test_inputs/train_w2v.csv'
7
+ with open('tests/test_inputs/train_w2v.csv', 'r') as f:
8
+ train_size = len(f.readlines()) - 1
9
+ assert dataset.__len__() == int(train_size * (1+dataset.neg_num))
10
+
@@ -0,0 +1,18 @@
1
+ import torch
2
+ from word2vec_trainer.word2vec_model import Word2Vec
3
+
4
+ def test_shape():
5
+ model = Word2Vec(10, 10)
6
+ input_embeddings = model.state_dict()['input_embedding_layer.weight']
7
+ output_embeddings = model.state_dict()['output_embedding_layer.weight']
8
+
9
+ assert input_embeddings.shape == (10, 10)
10
+ assert output_embeddings.shape == (10, 10)
11
+
12
+ def test_forward():
13
+ model = Word2Vec(10, 10)
14
+ x = torch.tensor([[1], [0], [4]])
15
+ y = torch.tensor([[2], [4], [8]])
16
+ z = model(x, y)
17
+
18
+ assert z.shape == (3, 10)
@@ -0,0 +1,20 @@
1
+ import os
2
+ from word2vec_trainer.preprocessing import Preprocessing
3
+
4
+ def test_preprocessing():
5
+ preprocessor = Preprocessing('tests/test_inputs/example_data.txt',
6
+ subsampled_file_path = 'tests/test_inputs/subsampled_data.txt',
7
+ vocabulary_file = 'tests/test_inputs/vocabulary.csv',
8
+ pos_example = 'tests/test_inputs/pos_examples.csv',
9
+ train_path = 'tests/test_inputs/train_w2v.csv',
10
+ valid_path = 'tests/test_inputs/valid_w2v.csv',
11
+ test_path = 'tests/test_inputs/test_w2v.csv')
12
+ preprocessor.run()
13
+
14
+ assert os.path.exists(preprocessor.train_path) == True
15
+ assert os.path.exists(preprocessor.valid_path) == True
16
+ assert os.path.exists(preprocessor.test_path) == True
17
+ assert os.path.exists(preprocessor.subsampled_file_path) == True
18
+ assert os.path.exists(preprocessor.vocabulary_file) == True
19
+ assert os.path.exists(preprocessor.pos_examples) == True
20
+
@@ -0,0 +1,16 @@
1
+ import os
2
+ from word2vec_trainer.trainer import Trainer
3
+
4
+ def test_trainer():
5
+ trainer = Trainer(
6
+ vocab_size = 10,
7
+ embedding_dim=2,
8
+ path_to_train='tests/test_inputs/train_w2v.csv',
9
+ path_to_val='tests/test_inputs/valid_w2v.csv',
10
+ path_to_neg_dist = 'tests/test_inputs/neg_distribution_test.npy',
11
+ epochs=10,
12
+ lr=0.01,
13
+ path_to_saved_model='tests/test_inputs/model.bin'
14
+ )
15
+
16
+ assert trainer.model != None