diff --git a/src/rbm/label.py b/src/rbm/label.py
index 427c30d..46c4407 100644
--- a/src/rbm/label.py
+++ b/src/rbm/label.py
@@ -3,11 +3,17 @@ from rbm.model import Model
 from rbm.entity import Entity, EntityParams, TrainingParams
 
 import math
+from enum import Enum
+
 class Label:
+	class EncodingType(Enum):
+		Binary = 1
+		OneHot = 2
+
 	class Fitter(Model):
 		def __init__(self, dim: tuple[int,int], work_dir: str = '.'):
 			super().__init__(f"label-{dim[0]}x{dim[1]}", work_dir)
-			self.unit1 = Entity(dim, EntityParams(num_gibbs_samples=1), TrainingParams(num_epochs=10000, do_rao_blackwell=False))
+			self.unit1 = Entity(dim, EntityParams(), TrainingParams(num_epochs=10000, do_rao_blackwell=False))
 
 		def forward(self, x: Mat) -> Mat:
 			return self.unit1.forward(x)
@@ -15,23 +21,16 @@ class Label:
 		def reconstruct(self, x: Mat):
 			return self.unit1.reconstruct(x)
 
-	def __init__(self, voc_size, num_dim, work_dir):
-		self.num_bits = round(math.log(voc_size,2))
-		self.fitter = Label.Fitter((self.num_bits, num_dim), work_dir)
-
-	def label2vec(self, list_of_labels: np.array):
-		batch = np.zeros((len(list_of_labels), self.num_bits))
-		for index, label in enumerate(list_of_labels):
-			b = self._enc_binary(label, self.num_bits)
-			batch[index, :] = np.array(b)
-
-		return batch
-
-	def vec2label(self, label_vecs: np.array, thresh=0.9):
-		labels = np.zeros((len(label_vecs)))
-		for index, label_vecs in enumerate(label_vecs):
-			labels[index] = self._dec_binary(label_vecs > thresh)
-		return labels
+	def __init__(self, voc_size, num_dim, encoding: EncodingType, work_dir):
+		self.voc_size = voc_size
+		if encoding == Label.EncodingType.Binary:
+			self.label2vec = self.label2vec_binary
+			self.vec2label = self.vec2label_binary
+			self.fitter = Label.Fitter((round(math.log(voc_size,2)), num_dim), work_dir)
+		elif encoding == Label.EncodingType.OneHot:
+			self.label2vec = self.label2vec_onehot
+			self.vec2label = self.vec2label_onehot
+			self.fitter = Label.Fitter((voc_size, num_dim), work_dir)
 
 	def fit(self, list_of_labels: np.array):
 		label_vecs = self.label2vec(list_of_labels)
@@ -46,6 +45,21 @@ class Label:
 		dec_data = self.fitter.reconstruct(list_of_encoded_label)
 		return dec_data
 
+	def label2vec_binary(self, list_of_labels: np.array):
+		num_bits = round(math.log(self.voc_size,2))
+		batch = np.zeros((len(list_of_labels), num_bits))
+		for index, label in enumerate(list_of_labels):
+			b = self._enc_binary(label, num_bits)
+			batch[index, :] = np.array(b)
+
+		return batch
+
+	def vec2label_binary(self, label_vecs: np.array, thresh=0.9):
+		labels = np.zeros((len(label_vecs)))
+		for index, label_vecs in enumerate(label_vecs):
+			labels[index] = self._dec_binary(label_vecs > thresh)
+		return labels
+
 	@staticmethod
 	def _enc_binary(number: int, num_bits):
 		res = [1 if s == '1' else 0 for s in format(number, f'{num_bits}b')]
@@ -61,15 +75,17 @@ class Label:
 
 		return res
 
-	@staticmethod
-	def num_dims(voc_size, num_codes_per_dim):
-		num_dim_required = 0
-		while voc_size > 1:
-			num_dim_required += 1
-			voc_size /= num_codes_per_dim
-
-		return num_dim_required
+	def label2vec_onehot(self, list_of_labels: np.array):
+		res = []
+		for label in list_of_labels:
+			vec = np.zeros(self.voc_size)
+			vec[label] = 1.0
+			res.append(vec)
+		return res
 
+	def vec2label_onehot(self, label_vecs: np.array, thresh=0.9):
+		# ToDo:
+		return None
 
 if __name__ == "__main__":
 	voc_size = 100