# cognitivess_model/tokenization_cognitivess.py

from transformers import PreTrainedTokenizer
import json

class CognitivessTokenizer(PreTrainedTokenizer):
    def __init__(self, vocab_file, merges_file=None, **kwargs):
        super().__init__(**kwargs)
        self.vocab_file = vocab_file
        self.merges_file = merges_file
        self.load_vocab()

    def load_vocab(self):
        # Load vocabulary
        with open(self.vocab_file, 'r') as f:
            self.vocab = {line.strip(): idx for idx, line in enumerate(f)}

        # Load merges file if exists
        self.merges = []
        if self.merges_file:
            with open(self.merges_file, 'r') as f:
                self.merges = [line.strip() for line in f]

    def _tokenize(self, text):
        # Tokenization logic (basic example)
        tokens = text.split()  # Simple whitespace-based tokenization
        return tokens

    def convert_tokens_to_ids(self, tokens):
        return [self.vocab.get(token, self.vocab.get('[UNK]')) for token in tokens]

    def convert_ids_to_tokens(self, ids):
        reverse_vocab = {idx: token for token, idx in self.vocab.items()}
        return [reverse_vocab.get(idx, '[UNK]') for idx in ids]

    def save_vocabulary(self, save_directory):
        vocab_path = f"{save_directory}/vocab.txt"
        with open(vocab_path, 'w') as f:
            for token in self.vocab:
                f.write(f"{token}\n")

        if self.merges_file:
            merges_path = f"{save_directory}/merges.txt"
            with open(merges_path, 'w') as f:
                for merge in self.merges:
                    f.write(f"{merge}\n")
            return vocab_path, merges_path
        return vocab_path,