Merge branch 'feature/raw-text-dataprep' of https://github.com/Vangmay/unsloth into feature/raw-text-dataprep
This commit is contained in:
vangmay
commit
d253b392fb
3 changed files with 143 additions and 132 deletions
|
|
@ -10,15 +10,16 @@ import tempfile
|
|||
from pathlib import Path
|
||||
import importlib.util
|
||||
|
||||
|
||||
# Mock the datasets module since it's not installed
|
||||
class MockDataset:
|
||||
def __init__(self, data_dict):
|
||||
self.data = data_dict
|
||||
self.column_names = list(data_dict.keys())
|
||||
|
||||
|
||||
def __len__(self):
|
||||
return len(next(iter(self.data.values())))
|
||||
|
||||
|
||||
def __getitem__(self, idx):
|
||||
if isinstance(idx, str):
|
||||
# Allow accessing columns by name like dataset['text']
|
||||
|
|
@ -28,19 +29,22 @@ class MockDataset:
|
|||
return {key: values[idx] for key, values in self.data.items()}
|
||||
else:
|
||||
raise TypeError(f"Invalid index type: {type(idx)}")
|
||||
|
||||
|
||||
@classmethod
|
||||
def from_dict(cls, data_dict):
|
||||
return cls(data_dict)
|
||||
|
||||
|
||||
# Mock datasets module
|
||||
datasets_mock = type(sys)('datasets')
|
||||
datasets_mock = type(sys)("datasets")
|
||||
datasets_mock.Dataset = MockDataset
|
||||
sys.modules['datasets'] = datasets_mock
|
||||
sys.modules["datasets"] = datasets_mock
|
||||
|
||||
# Import the raw_text module directly to avoid unsloth/__init__.py dependencies
|
||||
current_dir = os.path.dirname(__file__)
|
||||
raw_text_path = os.path.join(os.path.dirname(current_dir), 'unsloth', 'dataprep', 'raw_text.py')
|
||||
raw_text_path = os.path.join(
|
||||
os.path.dirname(current_dir), "unsloth", "dataprep", "raw_text.py"
|
||||
)
|
||||
|
||||
spec = importlib.util.spec_from_file_location("raw_text", raw_text_path)
|
||||
raw_text_module = importlib.util.module_from_spec(spec)
|
||||
|
|
@ -49,73 +53,75 @@ spec.loader.exec_module(raw_text_module)
|
|||
RawTextDataLoader = raw_text_module.RawTextDataLoader
|
||||
TextPreprocessor = raw_text_module.TextPreprocessor
|
||||
|
||||
|
||||
def test_raw_text_loader():
|
||||
"""Test basic RawTextDataLoader functionality."""
|
||||
|
||||
|
||||
# Mock tokenizer for testing
|
||||
class MockTokenizer:
|
||||
def __init__(self):
|
||||
self.eos_token = "</s>"
|
||||
|
||||
def __call__(self, text, return_tensors=None, add_special_tokens=False):
|
||||
|
||||
def __call__(self, text, return_tensors = None, add_special_tokens = False):
|
||||
words = text.split()
|
||||
token_ids = list(range(len(words)))
|
||||
|
||||
|
||||
if return_tensors == "pt":
|
||||
# Mock tensor-like object
|
||||
class MockTensor:
|
||||
def __init__(self, data):
|
||||
self.data = data
|
||||
|
||||
def __getitem__(self, idx):
|
||||
return self.data
|
||||
|
||||
def __len__(self):
|
||||
return len(self.data)
|
||||
|
||||
return {"input_ids": [MockTensor(token_ids)]}
|
||||
return {"input_ids": token_ids}
|
||||
|
||||
def decode(self, token_ids, skip_special_tokens=False):
|
||||
|
||||
def decode(self, token_ids, skip_special_tokens = False):
|
||||
return " ".join([f"word_{i}" for i in token_ids])
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
# Create test file
|
||||
test_content = "This is a test file for raw text training. " * 10
|
||||
with tempfile.NamedTemporaryFile(mode='w', suffix='.txt', delete=False) as f:
|
||||
with tempfile.NamedTemporaryFile(mode = "w", suffix = ".txt", delete = False) as f:
|
||||
f.write(test_content)
|
||||
test_file = f.name
|
||||
|
||||
|
||||
try:
|
||||
# Test loader
|
||||
tokenizer = MockTokenizer()
|
||||
loader = RawTextDataLoader(tokenizer, chunk_size=5, stride=2)
|
||||
|
||||
loader = RawTextDataLoader(tokenizer, chunk_size = 5, stride = 2)
|
||||
|
||||
# Test loading
|
||||
dataset = loader.load_from_file(test_file)
|
||||
assert len(dataset) > 0, "Should create at least one chunk"
|
||||
assert 'text' in dataset.column_names, "Dataset should have 'text' column"
|
||||
|
||||
assert "text" in dataset.column_names, "Dataset should have 'text' column"
|
||||
|
||||
# Test preprocessor
|
||||
preprocessor = TextPreprocessor()
|
||||
clean_text = preprocessor.clean_text(" messy text \n\n\n ")
|
||||
assert "messy text" in clean_text, "Should clean text properly"
|
||||
|
||||
|
||||
# Test validation
|
||||
stats = preprocessor.validate_dataset(dataset)
|
||||
assert stats['total_samples'] > 0, "Should count samples"
|
||||
assert 'warnings' in stats, "Should include warnings"
|
||||
|
||||
assert stats["total_samples"] > 0, "Should count samples"
|
||||
assert "warnings" in stats, "Should include warnings"
|
||||
|
||||
print("✅ All tests passed!")
|
||||
return True
|
||||
|
||||
|
||||
except Exception as e:
|
||||
print(f"❌ Test failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
finally:
|
||||
# Cleanup
|
||||
os.unlink(test_file)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
success = test_raw_text_loader()
|
||||
sys.exit(0 if success else 1)
|
||||
sys.exit(0 if success else 1)
|
||||
|
|
|
|||
|
|
@ -106,7 +106,7 @@ def run(args):
|
|||
# Use raw text loader
|
||||
loader = RawTextDataLoader(tokenizer, args.chunk_size, args.stride)
|
||||
dataset = loader.load_from_file(args.raw_text_file)
|
||||
elif args.dataset.endswith(('.txt', '.md', '.json', '.jsonl')):
|
||||
elif args.dataset.endswith((".txt", ".md", ".json", ".jsonl")):
|
||||
# Auto-detect local raw text files
|
||||
loader = RawTextDataLoader(tokenizer)
|
||||
dataset = loader.load_from_file(args.dataset)
|
||||
|
|
@ -409,35 +409,27 @@ if __name__ == "__main__":
|
|||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--raw_text_file",
|
||||
type=str,
|
||||
help="Path to raw text file for training"
|
||||
"--raw_text_file", type = str, help = "Path to raw text file for training"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--chunk_size",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Size of text chunks for training"
|
||||
"--chunk_size", type = int, default = 2048, help = "Size of text chunks for training"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stride",
|
||||
type=int,
|
||||
default=512,
|
||||
help="Overlap between chunks"
|
||||
"--stride", type = int, default = 512, help = "Overlap between chunks"
|
||||
)
|
||||
|
||||
TRAINING_MODES = {
|
||||
'instruction': 'Standard instruction-following',
|
||||
'causal': 'Causal language modeling (raw text)',
|
||||
'completion': 'Text completion tasks'
|
||||
"instruction": "Standard instruction-following",
|
||||
"causal": "Causal language modeling (raw text)",
|
||||
"completion": "Text completion tasks",
|
||||
}
|
||||
|
||||
parser.add_argument(
|
||||
"--training_mode",
|
||||
type=str,
|
||||
default="instruction",
|
||||
choices=list(TRAINING_MODES.keys()),
|
||||
help="Training mode for the model"
|
||||
type = str,
|
||||
default = "instruction",
|
||||
choices = list(TRAINING_MODES.keys()),
|
||||
help = "Training mode for the model",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
|
|
|||
|
|
@ -26,23 +26,24 @@ __all__ = [
|
|||
]
|
||||
|
||||
SUPPORTED_FORMATS = {
|
||||
'.txt': 'plain_text',
|
||||
'.md': 'markdown',
|
||||
'.json': 'json_lines',
|
||||
'.jsonl': 'json_lines',
|
||||
'.csv': 'csv_text_column'
|
||||
".txt": "plain_text",
|
||||
".md": "markdown",
|
||||
".json": "json_lines",
|
||||
".jsonl": "json_lines",
|
||||
".csv": "csv_text_column",
|
||||
}
|
||||
|
||||
|
||||
class RawTextDataLoader:
|
||||
def __init__(self, tokenizer, chunk_size=2048, stride=512):
|
||||
def __init__(self, tokenizer, chunk_size = 2048, stride = 512):
|
||||
self.tokenizer = tokenizer
|
||||
self.chunk_size = chunk_size
|
||||
self.chunk_size = chunk_size
|
||||
self.stride = stride
|
||||
|
||||
|
||||
def detect_format(self, file_path):
|
||||
"""Auto-detect file format and parse accordingly"""
|
||||
extension = Path(file_path).suffix.lower()
|
||||
return SUPPORTED_FORMATS.get(extension, 'plain_text')
|
||||
return SUPPORTED_FORMATS.get(extension, "plain_text")
|
||||
|
||||
def load_from_file(self, file_path):
|
||||
"""Load raw text and convert to dataset"""
|
||||
|
|
@ -50,7 +51,7 @@ class RawTextDataLoader:
|
|||
text_content = self._read_file_by_format(file_path, file_format)
|
||||
chunks = self.smart_chunk_text(text_content, self.chunk_size, self.stride)
|
||||
return self.create_causal_dataset(chunks)
|
||||
|
||||
|
||||
def load_from_files(self, file_paths):
|
||||
"""Load multiple text files"""
|
||||
all_chunks = []
|
||||
|
|
@ -61,11 +62,10 @@ class RawTextDataLoader:
|
|||
all_chunks.extend(chunks)
|
||||
return self.create_causal_dataset(all_chunks)
|
||||
|
||||
|
||||
def chunk_text(self, text):
|
||||
"""Split text into overlapping chunks"""
|
||||
return self.smart_chunk_text(text, self.chunk_size, self.stride)
|
||||
|
||||
|
||||
def create_causal_dataset(self, chunks):
|
||||
"""Create dataset for causal language modeling"""
|
||||
return Dataset.from_dict({"text": chunks})
|
||||
|
|
@ -79,51 +79,50 @@ class RawTextDataLoader:
|
|||
4. Adds proper EOS tokens
|
||||
"""
|
||||
# First pass: tokenize the entire text to get accurate token counts
|
||||
tokenized = self.tokenizer(text, return_tensors="pt", add_special_tokens=False)
|
||||
tokenized = self.tokenizer(text, return_tensors = "pt", add_special_tokens = False)
|
||||
tokens = tokenized["input_ids"]
|
||||
|
||||
|
||||
# Handle different tokenizer return formats
|
||||
if hasattr(tokens, '__len__') and len(tokens) > 0:
|
||||
if hasattr(tokens, "__len__") and len(tokens) > 0:
|
||||
# If it's a nested structure, get the first element
|
||||
if hasattr(tokens[0], '__len__'):
|
||||
if hasattr(tokens[0], "__len__"):
|
||||
tokens = tokens[0]
|
||||
elif isinstance(tokens, int):
|
||||
# If tokenizer returns just a count, create a simple range
|
||||
tokens = list(range(tokens))
|
||||
|
||||
|
||||
if len(tokens) <= chunk_size:
|
||||
# Text is small enough to fit in one chunk
|
||||
eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else ""
|
||||
return [text + eos_token]
|
||||
|
||||
|
||||
chunks = []
|
||||
start_idx = 0
|
||||
|
||||
|
||||
while start_idx < len(tokens):
|
||||
# Calculate end index for this chunk
|
||||
end_idx = min(start_idx + chunk_size, len(tokens))
|
||||
|
||||
|
||||
# Extract tokens for this chunk
|
||||
chunk_tokens = tokens[start_idx:end_idx]
|
||||
|
||||
|
||||
# Decode back to text
|
||||
chunk_text = self.tokenizer.decode(chunk_tokens, skip_special_tokens=True)
|
||||
|
||||
chunk_text = self.tokenizer.decode(chunk_tokens, skip_special_tokens = True)
|
||||
|
||||
# Add EOS token if it's the last chunk or chunk is complete
|
||||
if end_idx == len(tokens) or len(chunk_tokens) == chunk_size:
|
||||
eos_token = self.tokenizer.eos_token if self.tokenizer.eos_token else ""
|
||||
chunk_text += eos_token
|
||||
|
||||
|
||||
chunks.append(chunk_text)
|
||||
|
||||
|
||||
# Move to next chunk with stride overlap
|
||||
if end_idx == len(tokens):
|
||||
break
|
||||
start_idx += chunk_size - stride
|
||||
|
||||
|
||||
return chunks
|
||||
|
||||
|
||||
def tokenize_and_chunk(self, text):
|
||||
"""
|
||||
Tokenize first, then chunk by token count:
|
||||
|
|
@ -134,10 +133,10 @@ class RawTextDataLoader:
|
|||
|
||||
def _read_file_by_format(self, file_path, file_format):
|
||||
"""Read file content based on detected format."""
|
||||
with open(file_path, 'r', encoding='utf-8') as f:
|
||||
if file_format == 'plain_text' or file_format == 'markdown':
|
||||
with open(file_path, "r", encoding = "utf-8") as f:
|
||||
if file_format == "plain_text" or file_format == "markdown":
|
||||
return f.read()
|
||||
elif file_format == 'json_lines':
|
||||
elif file_format == "json_lines":
|
||||
lines = []
|
||||
for line in f:
|
||||
try:
|
||||
|
|
@ -147,42 +146,43 @@ class RawTextDataLoader:
|
|||
lines.append(text)
|
||||
except json.JSONDecodeError:
|
||||
continue
|
||||
return '\n\n'.join(lines)
|
||||
elif file_format == 'csv_text_column':
|
||||
return "\n\n".join(lines)
|
||||
elif file_format == "csv_text_column":
|
||||
reader = csv.DictReader(f)
|
||||
texts = []
|
||||
for row in reader:
|
||||
text = self._extract_text_from_csv_row(row)
|
||||
if text:
|
||||
texts.append(text)
|
||||
return '\n\n'.join(texts)
|
||||
return "\n\n".join(texts)
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_text_from_json(self, data):
|
||||
"""Extract text from JSON object using common field names."""
|
||||
text_fields = ['text', 'content', 'message', 'body', 'description', 'prompt']
|
||||
text_fields = ["text", "content", "message", "body", "description", "prompt"]
|
||||
for field in text_fields:
|
||||
if field in data and isinstance(data[field], str):
|
||||
return data[field]
|
||||
return ""
|
||||
|
||||
|
||||
def _extract_text_from_csv_row(self, row):
|
||||
"""Extract text from CSV row using common column names."""
|
||||
text_columns = ['text', 'content', 'message', 'body', 'description', 'prompt']
|
||||
text_columns = ["text", "content", "message", "body", "description", "prompt"]
|
||||
for column in text_columns:
|
||||
if column in row and row[column]:
|
||||
return row[column]
|
||||
return ""
|
||||
|
||||
|
||||
class TextPreprocessor:
|
||||
def clean_text(self, text):
|
||||
"""Remove unwanted characters, normalize whitespace"""
|
||||
text = re.sub(r'\s+', ' ', text)
|
||||
text = re.sub(r'[^\x20-\x7E\n\t]', '', text)
|
||||
text = text.replace('\r\n', '\n').replace('\r', '\n')
|
||||
text = re.sub(r'\n{3,}', '\n\n', text)
|
||||
text = re.sub(r"\s+", " ", text)
|
||||
text = re.sub(r"[^\x20-\x7E\n\t]", "", text)
|
||||
text = text.replace("\r\n", "\n").replace("\r", "\n")
|
||||
text = re.sub(r"\n{3,}", "\n\n", text)
|
||||
return text.strip()
|
||||
|
||||
|
||||
def extract_sections(self, text, patterns):
|
||||
"""Extract specific sections (e.g., code blocks, quotes)"""
|
||||
sections = []
|
||||
|
|
@ -193,12 +193,20 @@ class TextPreprocessor:
|
|||
|
||||
def add_structure_tokens(self, text):
|
||||
"""Add special tokens for structure (chapters, sections)"""
|
||||
text = re.sub(r'^# (.+)$', r'<|chapter|>\1<|/chapter|>', text, flags=re.MULTILINE)
|
||||
text = re.sub(r'^## (.+)$', r'<|section|>\1<|/section|>', text, flags=re.MULTILINE)
|
||||
text = re.sub(r'^### (.+)$', r'<|subsection|>\1<|/subsection|>', text, flags=re.MULTILINE)
|
||||
text = re.sub(r'```(\w*)\n(.*?)\n```', r'<|code|\1|>\2<|/code|>', text, flags=re.DOTALL)
|
||||
text = re.sub(
|
||||
r"^# (.+)$", r"<|chapter|>\1<|/chapter|>", text, flags = re.MULTILINE
|
||||
)
|
||||
text = re.sub(
|
||||
r"^## (.+)$", r"<|section|>\1<|/section|>", text, flags = re.MULTILINE
|
||||
)
|
||||
text = re.sub(
|
||||
r"^### (.+)$", r"<|subsection|>\1<|/subsection|>", text, flags = re.MULTILINE
|
||||
)
|
||||
text = re.sub(
|
||||
r"```(\w*)\n(.*?)\n```", r"<|code|\1|>\2<|/code|>", text, flags = re.DOTALL
|
||||
)
|
||||
return text
|
||||
|
||||
|
||||
def validate_dataset(self, dataset):
|
||||
"""
|
||||
Check for:
|
||||
|
|
@ -208,61 +216,66 @@ class TextPreprocessor:
|
|||
- Empty chunks
|
||||
"""
|
||||
stats = {
|
||||
'total_samples': len(dataset),
|
||||
'empty_samples': 0,
|
||||
'min_length': float('inf'),
|
||||
'max_length': 0,
|
||||
'avg_length': 0,
|
||||
'repeated_content': 0,
|
||||
'encoding_issues': 0,
|
||||
'warnings': []
|
||||
"total_samples": len(dataset),
|
||||
"empty_samples": 0,
|
||||
"min_length": float("inf"),
|
||||
"max_length": 0,
|
||||
"avg_length": 0,
|
||||
"repeated_content": 0,
|
||||
"encoding_issues": 0,
|
||||
"warnings": [],
|
||||
}
|
||||
|
||||
texts = dataset['text']
|
||||
|
||||
texts = dataset["text"]
|
||||
text_lengths = []
|
||||
seen_texts = set()
|
||||
|
||||
|
||||
for i, text in enumerate(texts):
|
||||
if not text or len(text.strip()) == 0:
|
||||
stats['empty_samples'] += 1
|
||||
stats["empty_samples"] += 1
|
||||
continue
|
||||
|
||||
|
||||
# Check for encoding issues
|
||||
try:
|
||||
text.encode('utf-8')
|
||||
text.encode("utf-8")
|
||||
except UnicodeEncodeError:
|
||||
stats['encoding_issues'] += 1
|
||||
|
||||
stats["encoding_issues"] += 1
|
||||
|
||||
# Calculate lengths
|
||||
length = len(text)
|
||||
text_lengths.append(length)
|
||||
stats['min_length'] = min(stats['min_length'], length)
|
||||
stats['max_length'] = max(stats['max_length'], length)
|
||||
|
||||
stats["min_length"] = min(stats["min_length"], length)
|
||||
stats["max_length"] = max(stats["max_length"], length)
|
||||
|
||||
# Check for repeated content
|
||||
text_hash = hash(text.strip())
|
||||
if text_hash in seen_texts:
|
||||
stats['repeated_content'] += 1
|
||||
stats["repeated_content"] += 1
|
||||
else:
|
||||
seen_texts.add(text_hash)
|
||||
|
||||
|
||||
# Calculate average length
|
||||
if text_lengths:
|
||||
stats['avg_length'] = sum(text_lengths) / len(text_lengths)
|
||||
stats['min_length'] = stats['min_length'] if stats['min_length'] != float('inf') else 0
|
||||
|
||||
# Generate warnings
|
||||
if stats['empty_samples'] > 0:
|
||||
stats['warnings'].append(f"Found {stats['empty_samples']} empty samples")
|
||||
|
||||
if stats['repeated_content'] > 0:
|
||||
stats['warnings'].append(f"Found {stats['repeated_content']} repeated samples")
|
||||
|
||||
if stats['encoding_issues'] > 0:
|
||||
stats['warnings'].append(f"Found {stats['encoding_issues']} encoding issues")
|
||||
|
||||
if stats['min_length'] < 10:
|
||||
stats['warnings'].append("Some samples are very short (< 10 characters)")
|
||||
|
||||
return stats
|
||||
stats["avg_length"] = sum(text_lengths) / len(text_lengths)
|
||||
stats["min_length"] = (
|
||||
stats["min_length"] if stats["min_length"] != float("inf") else 0
|
||||
)
|
||||
|
||||
# Generate warnings
|
||||
if stats["empty_samples"] > 0:
|
||||
stats["warnings"].append(f"Found {stats['empty_samples']} empty samples")
|
||||
|
||||
if stats["repeated_content"] > 0:
|
||||
stats["warnings"].append(
|
||||
f"Found {stats['repeated_content']} repeated samples"
|
||||
)
|
||||
|
||||
if stats["encoding_issues"] > 0:
|
||||
stats["warnings"].append(
|
||||
f"Found {stats['encoding_issues']} encoding issues"
|
||||
)
|
||||
|
||||
if stats["min_length"] < 10:
|
||||
stats["warnings"].append("Some samples are very short (< 10 characters)")
|
||||
|
||||
return stats
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue