from datasets import load_dataset

ds_train = load_dataset("ShoAnn/legalqa_klinik_hukumonline", split="train")
ds_test = load_dataset("ShoAnn/legalqa_klinik_hukumonline", split="test")

def gen(rows):
	sentences = []
	for question in rows["question"]:
		stripped = ''.join(c.lower() for c in question if c.isalpha() or c == ' ')
		sentences.append(stripped)
	return dict(sentence1=sentences)

ds_train = ds_train.map(gen, batched=True, remove_columns=ds_train.column_names)
ds_test = ds_test.map(gen, batched=True, remove_columns=ds_test.column_names)

class Model:
	def __init__(self):
		self.vocab = dict()

	@property
	def features(self):
		return list(self.vocab.keys())

	def train(self, ds):
		for row in ds:
			sentence = row["sentence1"]
			words = list(filter(lambda x: x, sentence.split()))

			for word in words:
				if word not in self.vocab:
					self.vocab[word] = 0
					print(f"Vocab size: {len(self.vocab)}", end='\r')
		print()

	def bag(self, sentence):
		counter = self.vocab.copy()

		# sanitation
		sanitized = ''.join(c for c in sentence if c.isalpha() or c == ' ')
		words = list(filter(lambda x: x, sanitized.split()))

		for word in words:
			if word not in counter:
				continue
			counter[word] = counter[word] + 1

		return list(counter.values())


# usage
PurpleBoW = Model()
PurpleBoW.train(ds_train)

excerpts = ds_test[:5]["sentence1"]
for e in excerpts:
	bag = PurpleBoW.bag(e)
	print(e, '->', bag)