AI_Summarizer / src /textSummarizer /components /data_transformation.py
Jeevant10's picture
deploy
1d2f1ad
Raw
History Blame Contribute Delete
2.67 kB
import os
from textSummarizer.logging import logger
from transformers import AutoTokenizer
from datasets import load_dataset, load_from_disk
from textSummarizer.entity import DataTransformationConfig
class DataTransformation:
def __init__(self, config: DataTransformationConfig):
self.config = config
# choose tokenizer checkpoint: prefer dev_model when dev_run is enabled
tokenizer_checkpoint = self.config.tokenizer_name
if getattr(self.config, 'dev_run', False) and getattr(self.config, 'dev_model', None):
tokenizer_checkpoint = self.config.dev_model
self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_checkpoint)
def convert_examples_to_features(self, example_batch):
input_encodings = self.tokenizer(
example_batch['dialogue'],
max_length=512,
truncation=True
)
target_encodings = self.tokenizer(
text_target=example_batch['summary'],
max_length=128,
truncation=True
)
return {
'input_ids': input_encodings['input_ids'],
'attention_mask': input_encodings['attention_mask'],
'labels': target_encodings['input_ids'],
}
def convert(self):
save_path = self.config.root_dir / "samsum_dataset"
# Temporarily disabled skip check to force transformation
# # Check if processed dataset already exists
# if save_path.exists() and (save_path / "dataset_dict.json").exists():
# logger.info(f"Processed dataset already exists at {save_path}. Skipping transformation.")
# return
logger.info(f"Loading dataset from {self.config.data_path}")
dataset_samsum = load_from_disk(str(self.config.data_path))
logger.info("Tokenizing dataset...")
dataset_samsum_pt = dataset_samsum.map(self.convert_examples_to_features, batched=True)
# ensure labels are correctly formatted for seq2seq Trainer
def rename_for_trainer(batch):
# already returns 'labels' in convert_examples_to_features; keep
return batch
dataset_samsum_pt = dataset_samsum_pt.map(rename_for_trainer, batched=True)
os.makedirs(save_path, exist_ok=True)
logger.info(f"Saving processed dataset to {save_path}")
# Use absolute path with proper Windows formatting to handle spaces in path
abs_save_path = os.path.abspath(str(save_path))
# Disable multiprocessing on Windows to avoid path issues with spaces
dataset_samsum_pt.save_to_disk(abs_save_path, num_proc=1)