|
|
| import sentencepiece as spm |
| import os |
| import json |
|
|
| class MTPTokenizer: |
| def __init__(self, model_path=None): |
| self.sp = None |
| self.model_path = model_path |
| if model_path and os.path.exists(model_path): |
| self.load(model_path) |
| |
| def train(self, corpus_path, vocab_size=4000, model_prefix='mtp_tokenizer'): |
| texts = [] |
| with open(corpus_path, 'r', encoding='utf-8') as f: |
| for line in f: |
| line = line.strip() |
| if not line: continue |
| try: |
| data = json.loads(line) |
| if 'instruction' in data: |
| texts.append(data['instruction']) |
| if 'input' in data and data['input'].strip(): |
| texts.append(data['input']) |
| if 'output' in data: |
| texts.append(data['output']) |
| except: continue |
|
|
| if not texts: raise ValueError("Corpus vacío") |
|
|
| temp_file = 'temp_corpus.txt' |
| with open(temp_file, 'w', encoding='utf-8') as f: |
| f.write('\n'.join(texts)) |
| |
| total_chars = sum(len(text) for text in texts) |
| min_vocab = 4000 |
| max_vocab = max(min_vocab, int(total_chars * 0.15)) |
| |
| try: |
| spm.SentencePieceTrainer.train( |
| input=temp_file, |
| model_prefix=model_prefix, |
| vocab_size=max_vocab, |
| model_type='bpe', |
| pad_id=0, unk_id=1, bos_id=2, eos_id=3, |
| character_coverage=1.0, |
| normalization_rule_name='identity', |
| num_threads=2, |
| split_digits=True, |
| max_sentencepiece_length=16 |
| ) |
| except RuntimeError as e: |
| if "Vocabulary size too high" in str(e): |
| import re |
| match = re.search(r'value <= (\d+)', str(e)) |
| if match: |
| spm.SentencePieceTrainer.train( |
| input=temp_file, |
| model_prefix=model_prefix, |
| vocab_size=int(match.group(1)), |
| model_type='bpe', |
| pad_id=0, unk_id=1, bos_id=2, eos_id=3, |
| character_coverage=1.0, |
| normalization_rule_name='identity', |
| num_threads=2 |
| ) |
| |
| os.remove(temp_file) |
| self.model_path = f"{model_prefix}.model" |
| self.load(self.model_path) |
| print(f"✓ Tokenizer trained: {self.vocab_size()} tokens") |
| |
| def load(self, model_path): |
| self.sp = spm.SentencePieceProcessor() |
| self.sp.load(model_path) |
| self.model_path = model_path |
| |
| def encode(self, text): |
| if self.sp is None: raise ValueError("Tokenizer not loaded") |
| return self.sp.encode_as_ids(text) |
| |
| def decode(self, ids): |
| if self.sp is None: raise ValueError("Tokenizer not loaded") |
| return self.sp.decode_ids(ids) |
| |
| def vocab_size(self): |
| if self.sp is None: return 0 |
| return self.sp.get_piece_size() |
| |
| def bos_id(self): return self.sp.bos_id() |
| def eos_id(self): return self.sp.eos_id() |
| def pad_id(self): return self.sp.pad_id() |
| def unk_id(self): return self.sp.unk_id() |
|
|