Spaces:
Sleeping
Sleeping
File size: 986 Bytes
5c7385e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 |
import re
from typing import List
class Splitter:
"""Base class for text splitting"""
def split(self, text: str) -> List[str]:
raise NotImplementedError
def join(self, tokens: List[str]) -> str:
raise NotImplementedError
class StringSplitter(Splitter):
"""Split text by pattern (default: space)"""
def __init__(self, split_pattern: str = ' '):
self.split_pattern = split_pattern
def split(self, prompt: str) -> List[str]:
return re.split(self.split_pattern, prompt.strip())
def join(self, tokens: List[str]) -> str:
return ' '.join(tokens)
class TokenizerSplitter(Splitter):
"""Split text using HuggingFace tokenizer"""
def __init__(self, tokenizer):
self.tokenizer = tokenizer
def split(self, prompt: str) -> List[str]:
return self.tokenizer.tokenize(prompt)
def join(self, tokens: List[str]) -> str:
return self.tokenizer.convert_tokens_to_string(tokens)
|